PS 문자열 알고리즘 - 6. 접미사 배열과 lcp 배열

 

이번에는 접미사 배열과 lcp 배열에 대해 살펴보고자 한다. 바로 본론으로 들어가자.

 

먼저, 접미사 배열은 어떤 문자열 $S$가 주어졌을 떄, 이 문자열의 모든 접미사를 사전순으로 정렬할 때, 각 접미사가 시작한 인덱스를 저장한 배열이다. 예를 들어서 문자열 ABRACADABRA를 살펴보자. 이 문자열의 모든 접미사는 아래와 같다.

 

ABRACADABRA 0
 BRACADABRA 1
  RACADABRA 2
   ACADABRA 3
    CADABRA 4
     ADABRA 5
      DABRA 6
       ABRA 7
        BRA 8
         RA 9
          A 10

 

이를 사전순으로 정렬하면 아래와 같은데, 이때 정렬된 상태에서 각 문자열이 원래 몇 번째 접미사였는지 각각 정수를 저장한 것을 접미사 배열(suffix array)라고 한다.

          A 10
       ABRA 7
ABRACADABRA 0
   ACADABRA 3
     ADABRA 5
        BRA 8
 BRACADABRA 1
    CADABRA 4
      DABRA 6
         RA 9
  RACADABRA 2

 

즉, ABRACADABRA의 접미사 배열은 $[10, 7, 0, 3, 5, 8, 1, 4, 6, 9, 2]$이다. 하지만, 이를 그대로 구현하면 모든 문자열의 접미사를 구해서 배열에 저장하는 것만 해도 $O(N^2)$이 소요된다. 하지만 이 $O(N^2)$은 너무 오래 걸리기 때문에 이를 줄이는 현명한 방법을 찾아야 한다.

 

어차피 실제로 문자열 전체를 가지고 정렬하지 않더라도, 가장 앞 자리만 봐도 이미 순서가 정해지는 것들이 있다. 예를 들면, BRA라는 접미사와 CADABRA라는 접미사는 가장 앞 글자만 비교해도 B와 C의 순서에 의해서 뒤쪽 접미사 전체를 확인하지 않고도 비교가 가능하다. 따라서 이런 방식을 활용하고자 하는 것이 접미사 배열을 빠르게 구하는 원리이다.

 

먼저 각 접미사의 첫 글자를 바탕으로 접미사 배열을 사용한다.

 

첫 글자만을 이용해서 정렬하면 A, B, C, D, R만 남는다. 그리고 여기에 0, 1, 2, 3, 4라는 인덱스를 붙이자.

그러면 이제 다시 모든 suffix들을 순회하면서 두 번째 글자를 확인하자.

 

지난 순회에서 매긴 인덱스를 이용해서 1st idx와 2nd idx의 값을 구할 수 있고, 그러면 길이 2의 튜플 형태로 쉽게 비교가 가능하고 새롭게 index를 매길 수 있다. (비어 있는 경우 가장 앞쪽으로 와야 하므로 -1을 매긴다.) 그리고 이 과정을 반복하는데, 이제 길이가 3인 suffix가 아닌 기존의 2배를 한 길이가 4인 suffix로 넘어가는 것이다.

 

이렇게 하면 지난 순회에서 길이 2의 문자열에 대해 각각 지난 순회에서 길이 2의 문자열에 대해 인덱스를 매겼으므로 길이 4의 문자열을 길이 2의 문자열 2개로 쪼개서 각각에 인덱스를 매길 수 있고, 그러면 정렬을 통해서 크기 2의 튜플의 비교로 정렬을 할 수 있다.

 

이를 반복하면 한 번에 문자열의 길이가 2배씩 증가하므로 문자열 길이가 $N$일 떄 대략 $log N$번 반복하고, 각 반복 과정에서 길이 $2$의 튜플을 $N$개 만들어 정렬하므로 $O(N log^2 N)$의 시간복잡도로 구할 수 있다.

 

그리고, 이 알고리즘의 시간복잡도를 log N을 제거할 수 있는데, 카운팅 정렬과 기수 정렬을 동시에 이용하는 것이다. key를 바탕으로 카운팅 정렬을 하는 함수를 먼저 구현하면 아래와 같다.

 

def count_sort(arr, key, k):
    # count sorting by key[idx], range: 0 - k
    n = len(arr)
    cnt = [0] * (k+1)
    for val in arr:
        cnt[key[val]] += 1
    pfs = [0] * (k+1)
    for i in range(k):
        pfs[i+1] = pfs[i] + cnt[i]
    out = [0] * n
    for val in arr:
        out[pfs[key[val]]] = val
        pfs[key[val]] += 1
    return out

 

 

arr라는 배열의 각각의 값 val을 key[val]의 내림차순으로 정렬하는 코드이다. 이때 key[val]의 값은 $0$ 이상 $k$ 이하여야 한다. 이 코드의 시간복잡도는 $O(n+k)$이지만, 우리가 suffix array를 만들 때는 $k$의 값이 항상 $n$을 넘지 않기 때문에 실제로는 $O(n)$이 된다.

 

길이 2인 튜플에 radix sort를 활용하면 count sort를 두 번 사용해서 정렬을 마무리지을 수 있고, 그러면 정렬의 시간복잡도를 $O(n)$으로 줄일 수 있다. 따라서 이를 활용해서 suffix array를 최종적으로 구하는 코드를 만들면 아래와 같다.

 

from bisect import bisect_left

def count_sort(arr, key, k):
    # count sorting by key[idx], range: 0 - k
    n = len(arr)
    cnt = [0] * (k+1)
    for val in arr:
        cnt[key[val]] += 1
    pfs = [0] * (k+1)
    for i in range(k):
        pfs[i+1] = pfs[i] + cnt[i]
    out = [0] * n
    for val in arr:
        out[pfs[key[val]]] = val
        pfs[key[val]] += 1
    return out

def make_suffix(s):
    n = len(s)
    if n == 0: return []
    sa = [*range(n)]
    uniq = []
    for ch in sorted(s):
        if not uniq or ch != uniq[-1]:
            uniq.append(ch)
    rank = [bisect_left(uniq, ch) for ch in s]
    k = 1
    while k < n:
        key1 = [r+1 for r in rank]
        key2 = [rank[i+k]+1 if i+k < n else 0 for i in range(n)]
        sa = count_sort(sa, key2, max(key2))
        sa = count_sort(sa, key1, max(key1))
        cnt = 0
        new_rank = [0] * n
        new_rank[sa[0]] = 0
        for i in range(1, n):
            a, b = sa[i], sa[i-1]
            ra, rb = rank[a], rank[b]
            rka = rank[a+k] if a+k < n else -1
            rkb = rank[b+k] if b+k < n else -1
            if ra != rb or rka != rkb:
                cnt += 1
            new_rank[a] = cnt
        rank = new_rank[:]
        if cnt == n-1: break
        k *= 2
    return sa

 

 

하지만 suffix array를 가지고 풀 수 있는 문제는 많지 않다. "기껏해야 suffix array를 구하라"라는 문제인 <BOJ 13264 접미사 배열 2> https://www.acmicpc.net/problem/13264 정도가 전부이고, 위 함수를 이용하면 되는 문제이니 위 함수를 직접 구현해보는 연습으로 삼아 풀어보면 될 것이다.

 

이 접미사 배열을 만든 이유는 lcp 배열을 만들기 위해서인데, lcp 배열은 Largest Common Prefix의 약자이고, suffix array에서 인접한 두 접미사끼리의 최대 공통 접두사의 길이가 저장된 배열이다. 즉, 위에서 abracadabra의 예시를 바탕으로 lcp 배열을 계산하면 아래와 같다.

                lcp
         A  10   X
       ABRA 7    1  (A, ABRA)
ABRACADABRA 0    4  (ABRA, ABRACADABRA)
   ACADABRA 3    1  (ABRACADABRA, ACADABRA)
     ADABRA 5    1  (ACADABRA, ADABRA)
        BRA 8    0  (ADABRA, BRA)
 BRACADABRA 1    3  (BRA, BRACADABRA)
    CADABRA 4    0  (BRACADABRA, CADABRA)
      DABRA 6    0  (CADABRA, DABRA)
         RA 9    0  (DABRA, RA)
  RACADABRA 2    2  (RA, RACADABRA)

 

다만, 맨 처음 인덱스는 이전 값이 없으므로 lcp 배열의 값이 존재하지 않는다. 따라서 lcp 배열의 경우 크기가 $n-1$이 되도록 만들 수도 있고, $lcp[0] = -1$이 되도록 구현하는 경우도 있다. 필자의 경우 lcp 배열의 크기가 $n-1$이 되게끔 구현하였다.

 

이를 구현하는 것도 그냥 구하면 시간복잡도가 최악의 경우 $O(N^2)$이 된다. 하지만 이를 $O(N)$으로 줄이는 방법이 있는데, 바로 lcp 배열을 suffix array 순이 아니라, 기존 문자열에서 앞쪽에서부터 시작한 접두사부터, 즉 길이가 긴 접미사부터 구하는 것이다. 따라서 원본 문자열의 각 접미사가 suffix array의 몇 번째 위치에 있었는지를 나타내는 배열 $rev$를 먼저 만들고, rev[0], rev[1], ... 과 같은 순으로 lcp를 계산하는 것이다.

 

ABRACADABRA에서 lcp 배열의 값을 계산하면 ABRA와 ABRACADABRA의 lcp 배열은 ABRA가 최장 공통 접두사이다. 그 다음으로 구하는 lcp 배열은 BRACADABRA의 lcp 값인데, ABRA라는 문자가 있었으므로 BRA라는 문자가 있음을 확신할 수 있고, 그러면 BRA와 BRACADABRA가 있음에 보장되므로 lcp 배열의 값은 최소 3이다. BRA와 BRACADABRA 사이에 어떤 접미사가 추가로 있을 수 있으므로 lcp 배열의 값이 3보다 클 수도 있다.

 

이렇게 되면 lcp 배열의 값의 최솟값이 보장되고, 직전의 lcp 값에서 최대 1만큼 감소하므로, 총 감소 횟수는 $N$이고, lcp 배열을 직접 증가시킬 때에는 그 인덱스가 전체 문자열에서 처음 인덱스 $0$에서 마지막 인덱스 $N-1$까지 이동하므로 증가 횟수는 감소 횟수 $N$에서 인덱스의 총 증가량 $N-1$을 더하면 $O(N)$에 작동함을 확인할 수 있다.

 

따라서 구현하면 아래와 같다. 이떄 $s$는 원본 문자열이고, $sa$는 위에서 구한 suffix array이다.

def makelcp(s, sa):
    n = len(s)
    if n <= 1: return []
    rev = [0] * n
    for i, p in enumerate(sa):
        rev[p] = i
    lcp = [0] * (n-1)
    prvl = 0
    for i in range(n):
        r = rev[i]
        if r == 0: prvl = 0; continue
        j = sa[r-1]
        while i+prvl < n and j+prvl < n and s[i+prvl] == s[j+prvl]:
            prvl += 1
        lcp[r-1] = prvl
        if prvl: prvl -= 1
    return lcp

 

 

 

이렇게 구한 lcp 배열을 이용해서 다양한 문제를 해결할 수 있다. 이는 다음 블로그에서 설명하려고 한다.