백준 9248번: Suffix Array 파이썬 (LCP 배열)

문제링크: 9248번: Suffix Array (acmicpc.net)

9248번: Suffix Array

이 문제는 단어의 Suffix Array를 구하고 LCP (Longest Common Prefix) 배열을 구하는 문제이다. Suffix Array란 단어의 접미사들을 알파벳 순서대로 정렬해 놓고, 정렬한 각 접미사들의 시작 인덱스를 저장한 배열이다. LCP 배열이란 Suffix Array를 구한 후 순서에 맞게 바로 이전 Suffix 와의 최장 공통 접두사 (LCP)의 길이를 저장한 배열이다. 이해가 어렵다면 문제 링크에 있는 banana 예시를 보면 된다.

접미사	인덱스	S.A	정렬된 접미사
banana	1	6	a
anana	2	4	ana
nana	3	2	anana
ana	4	1	banana
na	5	5	na
a	6	3	nana

Suffix Array를 구할때 모든 접미사를 앞글자부터 비교해 가며 정렬해 가면 시간 복잡도가 O(N^2)로 시간 초과가 난다. Suffix Array를 시간 복잡도 O(N logN) 에 구하는 방법은 다음과 같다. 각 접미사를 먼저 맨 앞 글자를 기준으로 정렬하여 그룹을 나눈다. 여기서 그룹이란 만약 첫번째 글자의 알파벳 순위가 같은 경우 (a, ana, anana) 에는 같은 그룹이 되는 것이다. 그러면 (a, ana, anana) 가 0그룹, (banana) 가 1그룹, (na, nana) 가 2그룹이 된다. 다음으로 위 그룹 내에서 맨 앞글자에서 변화하는 간격 d=1 만큼 인덱스를 증가시켜 비교하고 새 그룹이 생긴다면 그룹을 늘려주어 새로 할당한다. 그러면 (a) 0그룹, (ana, anana) 1그룹, (banana) 2그룹, (na, nana) 3그룹이 된다. 다음으로 간격 d에 2씩 곱해주면서 비교를 모든 접미사가 N-1 개  (0그룹부터 시작했으므로) 의 그룹으로 나뉠때 까지 반복한다.

 

위 과정에서 d는 최대 logN, 각 단어를 비교하는 과정은 N 이므로 시간 복잡도 O(N logN) 안에 Suffix Array를 만들 수 있게 된다. 그러면 이론적인 원리는 알았는데, 구현을 어떻게 할지 도무지 감이 안잡히게 된다. 힌트는 먼저 그룹을 나누어 주고, 그룹을 기준으로 배열을 정렬하면 각 d에 대한 S.A.를 얻을 수 있다. 

import sys

word = sys.stdin.readline().rstrip()
N = len(word)
sa = [-1 for _ in range(N)] #sa[i] - i번째 suffix array 단어의 시작인덱스
pos = [0 for _ in range(N)] #pos[x] - 시작인덱스 x인 suffix의 그룹

def get_pos(x,d): #return pos[x+d]
    if x+d<N: return pos[x+d]
    else: return -1

def temp_same(x,y): #x,y d*=2 후의 그룹할당
    if pos[x]<pos[y]: return 1
    if pos[x]==pos[y]:
        if get_pos(x,d)<get_pos(y,d): return 1
    return 0

for i in range(N):
    sa[i] = i
    pos[i] = ord(word[i])-96


temp = [0 for _ in range(N)]
sa = sorted(sa, key=lambda x: (pos[x]))

d = 1
while True:
    sa = sorted(sa, key=lambda x: (pos[x], get_pos(x,d)))
    for i in range(N-1):
        temp[i+1] = temp[i] + temp_same(sa[i], sa[i+1])
    for i in range(N):
        pos[sa[i]] = temp[i]
    if temp[N-1]==N-1:
        break
    d*=2

order = [-1 for _ in range(N)]
for i in range(N):
    order[sa[i]] = i
lcp = [-1 for _ in range(N)]

d=0
for i in range(N):
    k = order[i]
    if k:
        j = sa[k-1]
        while j+d < N and i+d < N and word[j+d]==word[i+d]:
            d+=1
        lcp[k]=d
        if d:
            d -=1
    
for x in sa:
    print(x+1, end=' ')
print()
for x in lcp:
    if x == -1:
        print('x', end=' ')
    else:
        print(x, end=' ')

먼저 sa[i] 리스트는 단어의 Suffix Array를 저장하고, pos[i] 리스트는 i번째 접미사의 그룹을 저장하고, temp[i]는 정렬된 Suffix Array의 그룹을 저장한다. 만약 Suffix Array가 완성되었다면 그룹이 총 N-1개 있다는 뜻이므로 temp[N-1] == N-1이 될 것이다.

 

get_pos(x,d) 함수는 현재 인덱스에서 d만큼 증가한 그룹 pos[x+d]를 반환한다. 만약 접미사 x+d번째 인덱스가 공백이라면 그 어떤 알파벳보다 우선 순위가 가장 높아져야 하므로 -1을 반환하도록 했다. temp_same(x,y) 함수는 같은 그룹 내의 x와 y의 d에 따라 그룹이 새로 나뉘는지, 나뉘지 않는지를 판단하는 함수이다. 그룹이 새로 나뉘면 1, 나뉘지 않으면 0을 반환하여 새로 나뉘는 그룹이 있다면 temp[i+1] = temp[i] + 1이 되어 그룹의 개수를 표현하는 temp가 1씩 증가하고 아니면 유지된다. 

 

먼저 pos[i] 리스트를 각 i번째 접미사의 첫 글자로 그룹을 할당하고, pos[i]를 key로 sa리스트를 정렬시켜준다. 다음 d에 따라 key를 (pos[x], get_pos(x,d))를 사용해 현재 그룹으로 정렬 + pos[x+d] 를 고려한 정렬을 하여 d에 대한 정렬을 시켜준다. 이후 S.A. 를 앞부터 훑어 가며 sa[i]와 sa[i+1]을 당시의 d를 사용해 비교하여 만약 당시의 d 비교로 그룹이 나뉜다면 temp 배열을 통해 새로운 그룹을 할당한다. temp[i]는 현재 그룹의 개수를 파악하기 위해 만든 배열로 만약 모든 접미사가 모든 그룹으로 나뉘었다면 temp[N-1]=N-1 이 될 것이다. 

 

이후 Suffix Array를 통해 LCP를 구하는 방법은 간단하다. order[i]=k는 i번째 접미사의 Suffix Array의 인덱스를 저장하고 있다. sa[k]는 i번째 접미사의 Suffix Array 숫자이고, sa[k-1]는 i번째 접미사의 Suffix Array로 정렬후 바로 전 접미사이다. banana의 경우 i부터 훑은 word[sa[order[i]]:] = word[i:] , word[sa[order[i]-1]:] = word[j:] 와 비교한 각 LCP는 다음과 같다.

 

banana anana / 0
anana ana / 3
nana na / 2
ana a / 1
na banana / 0

 

보면 order = i 의 word[i:]와 word[j:]의 LCP가 1보다 크면, 다음 order = i+1 의 word[i+1:]와 word[sa[order[i+1]-1]:] 의 LCP는 order = i의 LCP에 1을 뺀 값보다 큰 것을 알 수 있다.  예시로 LCP가 3이라면 word[i:]와 word[j:] 의 앞 세글자가 같다는 뜻인데, word[i:]의 접미사와 word[j:]의 접미사 또한 접미사 배열에 존재하므로 word[i+1:]와 word[j:][1:] 의 LCP는 무조건 3-1 = 2 일 것이기 때문이다. 위 규칙을 이용하면 LCP배열을 O(N)의 시간 복잡도를 구할 수 있다.