문자열 핵심 자료구조 suffix array O(NlogN) 최적화 배우기

1. 최적화

 

이전 문자열 자료구조 suffix array를 만드는 manber-myers 알고리즘은 $O(Nlog^{2}N)$의 시간복잡도를 가진다.

 

https://deepdata.tistory.com/967 

문자열 핵심 자료구조 suffix array와 lcp(longest common prefix)배열 구하는법 배우기(Manber-Myers, kasai)

 

def get_suffix_array(s):
    
    n = len(s)

    #i+d/2번째 접미사의 rank를 찾다가, 넘어가는 경우 빈 문자열의 rank는 0으로 하기 위해
    rank = [[0,0,0] for _ in range(n)]

    #d = 1일때 정렬단계
    for i in range(n):

        #첫번째 원소는 i번째 접미사의 rank
        #i번째 접미사의 첫 문자를 아스키코드로 바꿔놓은 값으로 바꿔놓는다
        #소문자로만 이루어져있다고 가정

        #두번째 원소는 i+d/2번째 접미사의 rank

        #세번째 원소(suffix array)는 i번째 접미사의 시작 인덱스 i로 초기화
        rank[i] = [ord(s[i]) - ord('a') + 1,0,i]
    
    
    d = 1 #d/2를 나타내는 값

    while d < 2*n:
        
        #d단계 정렬 O(NlogN)정렬
        #첫번째, 두번째 원소로만 정렬하도록
        rank = sorted(rank,key = lambda x: [x[0],x[1]])

        #rank 재구성

        new_rank = [0]*n
        new_rank[0] = 1

        for i in range(1,n):
            
            if rank[i-1][0] == rank[i][0] and rank[i-1][1] == rank[i][1]:
                
                new_rank[i] = new_rank[i-1]
            
            else:
                
                new_rank[i] = new_rank[i-1] + 1
        
        #모든 접미사가 1~n개의 그룹으로 나뉘면 접미사 배열이 완성된 것
        #새롭게 부여된 rank는 첫번째 원소를 구성하는데,
        #모두 서로 다르다면 순서쌍을 정렬 하더라도 변화 없기 때문이다.
        if new_rank[n-1] == n:
            
            break
        
        #i번째 접미사의 rank를 저장
        ind_rank = [0]*n

        for i in range(n):
            
            rank[i][0] = new_rank[i]
            ind_rank[rank[i][2]] = rank[i][0]

        #순서쌍 재구성
        #두번째 원소를 i+d/2번째 접미사의 rank로 구성
        for i in range(n):
            
            #i번째 접미사의 위치 i = rank[i][2]
            #d/2칸 앞의 위치는 rank[i][2] + d

            #d/2칸 앞 위치가 길이를 넘어가면, 0으로
            if rank[i][2]+d >= n:
                
                rank[i][1] = 0

            #길이를 넘어가지 않으면...
            #i+d/2번째 접미사의 rank는 ind_rank[rank[i][2]+d]에 있다

            else:
                
                rank[i][1] = ind_rank[rank[i][2]+d]

        #다음 d 값을 2배 하기
        d *= 2
    
    suffix = [0]*(n)

    for i in range(n):
        
        suffix[i] = rank[i][2] + 1

    return suffix

O(logN)번 O(NlogN) 정렬을 수행하기 때문에 $O(Nlog^{2}N)$정도 걸리는데

 

#d단계 정렬 O(NlogN)정렬
#첫번째, 두번째 원소로만 정렬하도록
rank = sorted(rank,key = lambda x: [x[0],x[1]])

 

counting sort를 이용한 O(N) 정렬을 수행하면 O(NlogN)으로 줄일 수 있다

 

tuple형태의 rank 배열을 첫번째 원소, 두번째 원소 기준으로 정렬하기 떄문에

 

counting sort를 하려면 stable sort 방법을 알아야한다

 

https://deepdata.tistory.com/968

파이썬 pair type stable sort 배우기(tuple counting sort)

 

counting sort 함수를 직접 만들고 정렬하는 부분만 counting sort로 바꾸면 된다.

 

def counting_sort(rank,n):

    #rank 배열의 크기(문자열의 길이): n
    
    #가능한 문자의 종류 수(알파벳 소문자를 가정)
    #첫 단계에는 rank 수가 26을 넘을 수 있다
    m = max(26,n)

    #먼저 두번째 원소를 기준으로 정렬함
    counting = [0]*(m+1)

    for i in range(n):
    
        counting[rank[i][1]] += 1

    for i in range(1,m+1):
    
        counting[i] += counting[i-1]
    
    #두번째 원소를 기준으로 정렬한 인덱스를 저장
    second = [0]*n

    for i in range(n-1,-1,-1):
    
        second[counting[rank[i][1]]-1] = i
    
        counting[rank[i][1]] -= 1
    
    #첫번째 원소를 기준으로 정렬
    counting = [0]*(m+1)

    for i in range(n):
    
        counting[rank[i][0]] += 1

    for i in range(1,m+1):
    
        counting[i] += counting[i-1]
    
    #정렬된 원소를 담을 rank배열
    new_rank = [0]*(n)

    for i in range(n-1,-1,-1):
    
        new_rank[counting[rank[second[i]][0]]-1] = rank[second[i]]

        counting[rank[second[i]][0]] -= 1
    
    return new_rank

def get_suffix_array(s):

    n = len(s)

    #i+d/2번째 접미사의 rank를 찾다가, 넘어가는 경우 빈 문자열의 rank는 0으로 하기 위해
    rank = [[0,0,0] for _ in range(n)]

    #d = 1일때 정렬단계
    for i in range(n):

        #첫번째 원소는 i번째 접미사의 rank
        #i번째 접미사의 첫 문자를 아스키코드로 바꿔놓은 값으로 바꿔놓는다
        #소문자로만 이루어져있다고 가정

        #두번째 원소는 i+d/2번째 접미사의 rank

        #세번째 원소(suffix array)는 i번째 접미사의 시작 인덱스 i로 초기화
        rank[i] = [ord(s[i]) - ord('a') + 1,0,i]


    d = 1 #d/2를 나타내는 값

    while d < 2*n:

        #d단계 정렬 O(N) counting sort
        rank = counting_sort(rank,n)

        #rank 재구성

        new_rank = [0]*n
        new_rank[0] = 1

        for i in range(1,n):

            if rank[i-1][0] == rank[i][0] and rank[i-1][1] == rank[i][1]:

                new_rank[i] = new_rank[i-1]

            else:

                new_rank[i] = new_rank[i-1] + 1

        #모든 접미사가 1~n개의 그룹으로 나뉘면 접미사 배열이 완성된 것
        #새롭게 부여된 rank는 첫번째 원소를 구성하는데,
        #모두 서로 다르다면 순서쌍을 정렬 하더라도 변화 없기 때문이다.
        if new_rank[n-1] == n:

            break

        #i번째 접미사의 rank를 저장
        ind_rank = [0]*n

        for i in range(n):

            rank[i][0] = new_rank[i]
            ind_rank[rank[i][2]] = rank[i][0]

        #순서쌍 재구성
        #두번째 원소를 i+d/2번째 접미사의 rank로 구성
        for i in range(n):

            #i번째 접미사의 위치 i = rank[i][2]
            #d/2칸 앞의 위치는 rank[i][2] + d

            #d/2칸 앞 위치가 길이를 넘어가면, 0으로
            if rank[i][2]+d >= n:

                rank[i][1] = 0

            #길이를 넘어가지 않으면...
            #i+d/2번째 접미사의 rank는 ind_rank[rank[i][2]+d]에 있다

            else:

                rank[i][1] = ind_rank[rank[i][2]+d]

        #다음 d 값을 2배 하기
        d *= 2

    suffix = [0]*(n)

    for i in range(n):

        suffix[i] = rank[i][2] + 1

    return suffix

def kasai(suffix_array, s):

    n = len(s)

    lcp = [0]*n

    #suffix_array[i] = 사전순으로 i번째 원소가 몇번째 접미사인지.
    #reverse_suffix_array[i] = i번째 접미사가 사전 순으로 몇번째 원소인지

    #suffix_array[k] = i <<>> reverse_suffix_array[i] = k

    #i번째 접미사 suffix_array[k] = i에서 앞 한글자를 지운
    #i+1번째 접미사 suffix_array[?] = i+1의 ?를 알아야함

    reverse_suffix_array = [0]*n

    for i in range(n):

        #suffix_array[i] = k <<<>>> reverse_suffix_array[k] = i
        reverse_suffix_array[suffix_array[i] - 1] = i

    #lcp 배열을 구한다

    #가장 긴 접미사인 0번째 접미사부터 시작

    #0번째 접미사는 reverse_suffix_array[0] = i위치에 존재하고
    #0번째 접미사와 인접한 접미사는... suffix_array[i-1] = suffix_array[reverse_suffix_array[0]-1]에 존재
    #두 접미사를 비교해서 몇개가 일치하는지 구한다

    k = 0 #lcp값

    for i in range(n):

        #lcp[0]는 정의하지 않는다
        #suffix_array의 0번째 원소의 접미사는 이전 접미사가 존재하지 않는다
        if reverse_suffix_array[i] == 0:

            lcp[0] = 'x'
            k = 0
            continue

        #i번째 접미사와 인접한 접미사의 위치
        adj = suffix_array[reverse_suffix_array[i]-1] - 1

        #i번째 접미사와 인접한 접미사의 글자를 하나씩 비교하면서
        #lcp값 k를 1씩 증가시킨다.
        while i+k < n and adj+k < n and s[i+k] == s[adj+k]:

            k += 1

        #반복문을 탈출하면, i번째 접미사에 대응하는 lcp값은 k가 된다.
        lcp[reverse_suffix_array[i]] = k

        # 그 다음 반복문에서 앞 1글자를 지우므로, lcp는 적어도 1 감소해야한다.
        if k > 0:

            k -= 1

    return lcp

s = input()

suffix = get_suffix_array(s)
lcp = kasai(suffix,s)

print(*suffix)
print(*lcp)

 

여기서 주의할 부분은 counting 배열의 class 수가 알파벳 소문자를 가정한다면 무조건 m = 26이라고 생각하기 쉬운데

 

rank배열의 순서쌍 (a,b)를 기준으로 class를 부여하기 때문에 26을 넘을 수 있다..

 

(1,1), (1,2), (1,3), ... , (1,26), (2,1),.... 

 

이런식으로 서로 다른 순서쌍이 26개 넘으면 새로운 rank 부여할때 

 

[1,2,3,...,26,27,...]으로 26을 넘어가거든

 

azfhazfgifjpfoqjfpfofjjfqpowfjefpdqkfkfdj

[[1, 0, 0], [26, 0, 1], [6, 0, 2], [8, 0, 3], [1, 0, 4], [26, 0, 5], [6, 0, 6], [7, 0, 7], [9, 0, 8], [6, 0, 9], [10, 0, 10], [16, 0, 11], [6, 0, 12], [15, 0, 13], [17, 0, 14], [10, 0, 15], [6, 0, 16], [16, 0, 17], [6, 0, 18], [15, 0, 19], [6, 0, 20], [10, 0, 21], [10, 0, 22], [6, 0, 23], [17, 0, 24], [16, 0, 25], [15, 0, 26], [23, 0, 27], [6, 0, 28], [10, 0, 29], [5, 0, 30], [6, 0, 31], [16, 0, 32], [4, 0, 33], [17, 0, 34], [11, 0, 35], [6, 0, 36], [11, 0, 37], [6, 0, 38], [4, 0, 39], [10, 0, 40]]
[[1, 0, 0], [26, 0, 1], [6, 0, 2], [8, 0, 3], [1, 0, 4], [26, 0, 5], [6, 0, 6], [7, 0, 7], [9, 0, 8], [6, 0, 9], [10, 0, 10], [16, 0, 11], [6, 0, 12], [15, 0, 13], [17, 0, 14], [10, 0, 15], [6, 0, 16], [16, 0, 17], [6, 0, 18], [15, 0, 19], [6, 0, 20], [10, 0, 21], [10, 0, 22], [6, 0, 23], [17, 0, 24], [16, 0, 25], [15, 0, 26], [23, 0, 27], [6, 0, 28], [10, 0, 29], [5, 0, 30], [6, 0, 31], [16, 0, 32], [4, 0, 33], [17, 0, 34], [11, 0, 35], [6, 0, 36], [11, 0, 37], [6, 0, 38], [4, 0, 39], [10, 0, 40]]
[[1, 14, 0], [1, 14, 4], [2, 12, 33], [2, 8, 39], [3, 4, 30], [4, 6, 2], [4, 5, 6], [4, 8, 9], [4, 10, 12], [4, 11, 16], [4, 10, 18], [4, 8, 20], [4, 12, 23], [4, 8, 28], [4, 11, 31], [4, 9, 36], [4, 2, 38], [5, 7, 7], [6, 1, 3], [7, 4, 8], [8, 11, 10], [8, 4, 15], [8, 8, 21], [8, 4, 22], [8, 3, 29], [8, 0, 40], [9, 4, 35], [9, 4, 37], [10, 12, 13], [10, 4, 19], [10, 13, 26], [11, 4, 11], [11, 4, 17], [11, 10, 25], [11, 2, 32], [12, 8, 14], [12, 11, 24], [12, 9, 34], [13, 4, 27], [14, 4, 1], [14, 4, 5]]
[[1, 7, 0], [1, 6, 4], [2, 0, 39], [3, 21, 33], [4, 25, 30], [5, 16, 38], [6, 15, 6], [7, 1, 2], [8, 26, 9], [8, 18, 20], [8, 4, 28], [9, 5, 36], [10, 28, 12], [10, 8, 18], [11, 10, 16], [11, 3, 31], [12, 27, 23], [13, 8, 7], [14, 32, 3], [15, 20, 8], [16, 0, 40], [17, 11, 29], [18, 26, 15], [18, 30, 22], [19, 12, 21], [20, 10, 10], [21, 21, 35], [21, 2, 37], [22, 19, 19], [23, 18, 13], [24, 8, 26], [25, 29, 32], [26, 23, 11], [26, 22, 17], [27, 31, 25], [28, 11, 14], [29, 9, 34], [30, 24, 24], [31, 17, 27], [32, 14, 1], [32, 13, 5]]

 

 

2. 다른 사람 코드로 최적화하기

 

저렇게 하면 쓸데없는걸 많이 해서 그런가..

 

이전에 $O(Nlog^{2}N)$이 6초정도인데 $O(NlogN)$이면 5초정도로 1초는 줄어들긴 하더라

 

다른 사람 코드 보면 1초 정도에 깔끔하게 구현하더라고... 그걸로 바꾸자

 

완전히 최적화된 다음은 0.5초 정도에 통과한다..?

 

def counting_sort(suffix_array,rank,d):
    
    m = max(n,256) #가질 수 있는 rank 수
    
    counting = [0]*(m+1)

    #두번째 원소인 i+d번째 rank에 대한 정렬 단계
    #rank의 두번째 원소는 d번부터

    #i번째 접미사와 i+d번째 접미사의 rank차이가 d이기 때문
    #계산된 rank 누적합에 d를 더해줘야함
    counting[0] = d 

    for i in range(d,n):
        
        counting[rank[i]] += 1
    
    for i in range(1,m+1):
        
        counting[i] += counting[i-1]
    
    second = [0]*n

    for i in range(n-1,-1,-1):
        
        if i+d >= n: #index가 넘어가는 경우, 추가로 만든 n번 index로 
            
            ind = n
        
        else:
            
            ind = i+d
        
        counting[rank[ind]] -= 1
        second[counting[rank[ind]]] = i
    
    #첫번째 원소에 대한 정렬

    counting = [0]*(m+1)

    for i in range(n):
        
        counting[rank[i]] += 1
    
    for i in range(1,m+1):
        
        counting[i] += counting[i-1]
    
    for i in range(n-1,-1,-1):
        
        counting[rank[second[i]]] -= 1

        suffix_array[counting[rank[second[i]]]] = second[i]
    
    return suffix_array

def compare(rank,a,b,d):
    
    if rank[a] == rank[b] and rank[a+d] == rank[b+d]:

        return 0

    else:

        return 1 

def get_suffix_array(suffix_array,rank):
    
    d = 1

    while d < 2*n:
        
        suffix_array = counting_sort(suffix_array,rank,d)

        new_rank = [0]*(n+1)

        new_rank[suffix_array[0]] = 1

        for i in range(1,n):
            
            new_rank[suffix_array[i]] = new_rank[suffix_array[i-1]] + compare(rank,suffix_array[i-1],suffix_array[i],d)
        
        #모든 rank가 서로 다르다면 정렬이 끝났다
        if new_rank[suffix_array[n-1]] == n:
            
            break

        rank = new_rank

        d *= 2
    
    return suffix_array

def kasai(s, suffix_array):
    
    n = len(s)

    lcp = [0]*n

    reverse_suffix_array = [0]*n

    for i in range(n):
        
        reverse_suffix_array[suffix_array[i]] = i
    
    k = 0

    for i in range(n):
        
        if reverse_suffix_array[i] == 0:
            
            lcp[0] = 'x'
            k = 0
            continue
        
        adj = suffix_array[reverse_suffix_array[i]-1]

        while i+k < n and adj+k < n and s[i+k] == s[adj+k]:
            
            k += 1
        
        lcp[reverse_suffix_array[i]] = k

        if k > 0:
            
            k -= 1
    
    return lcp

s = input()
n = len(s)

#초기화 단계
suffix_array = [i for i in range(n)]

#일차원 배열 rank에 첫번째 원소, 두번째 원소를 모두 넣을거임
#i번째 접미사의 첫번째 원소는 i번째에, 두번째 원소는 i+d번째에 있겠지
#i+d가 넘어가면 0으로 하면 되고

rank = [ord(s[i])-ord('a')+1 for i in range(n)]

rank.append(0) #i+d가 넘어가는 경우

suffix_array = get_suffix_array(suffix_array,rank)
lcp = kasai(s,suffix_array)

#배열에 대한 빠른 출력
print(" ".join(map(lambda i:str(i+1), suffix_array)))
print(" ".join(map(str,lcp)))

 

2-1) rank 배열이 3차원이어야하는가?

 

핵심은 나는 rank배열을 3차원 배열로 만들었지만.. 굳이 그럴필요 있는가에 대한 의문이다

 

0~n-1의 rank배열에 대해 i번째 접미사의 첫번째 rank는 rank[i]에 있고 두번째 rank는?

 

두번째 rank는 i+d번째 접미사의 rank를 저장하는데 사실 rank[i+d]에 있다는 것이다.

 

그래서 굳이 3차원 배열로 rank를 만들 필요 없이 1차원 배열의 rank를 선언

 

s = input()
n = len(s)

#초기화 단계
suffix_array = [i for i in range(n)]

#일차원 배열 rank에 첫번째 원소, 두번째 원소를 모두 넣을거임
#i번째 접미사의 첫번째 원소는 i번째에, 두번째 원소는 i+d번째에 있겠지
#i+d가 넘어가면 0으로 하면 되고

rank = [ord(s[i])-ord('a')+1 for i in range(n)]

rank.append(0) #i+d가 넘어가는 경우

 

여기서 i+d가 n-1을 넘어갈 수 있는데 그런 경우에는 0을 부여하기로 했다.

 

그러면 d< 2n까지 반복해야하니까 rank 크기를 2n정도로 잡아야하나???

 

그럴 필요 없이 i+d가 n-1을 넘어가면 rank[n]을 부여하면 된다.

 

그러므로 n번째 원소에 0 하나만을 추가해놓는다.

 

 

2-2) suffix_array에 대한 의문

 

그리고 suffix_array도 rank와 따로 초기화해놓는데...

 

나는 rank의 3번째 원소에 suffix_array 원소를 저장해놓았는데.. 이게 사실 그럴 필요가 있는가를 잘 생각해봐야한다

 

잘 생각해보면.. "rank배열의 정렬된 결과 인덱스를 담은 배열이 suffix_array가 된다"

 

그러니까 rank배열을 counting sort하고 sorting된 index가 suffix_array면 된다는 것이다.

 

 

 

2-3) 1차원 배열로 tuple counting sort하기

 

두번째 원소를 기준으로 정렬하고, 첫번째 원소를 기준으로 정렬한다.

 

rank배열의 두번째 원소는 i+d번째 접미사의 rank값이다.

 

i = 0번째 접미사의 i+d번째 접미사는 d번이기 때문에 d번부터 n-1번까지 순회하여 counting배열을 구성하면..

 

두번째 원소를 기준으로 counting배열을 구성하게 된다.

 

def counting_sort(suffix_array,rank,d):
    
    m = max(n,256) #가질 수 있는 rank 수
    
    counting = [0]*(m+1)

    #두번째 원소인 i+d번째 rank에 대한 정렬 단계
    #rank의 두번째 원소는 d번부터

    #i번째 접미사와 i+d번째 접미사의 rank차이가 d이기 때문
    #계산된 rank 누적합에 d를 더해줘야함
    counting[0] = d 

    for i in range(d,n):
        
        counting[rank[i]] += 1
    
    for i in range(1,m+1):
        
        counting[i] += counting[i-1]

 

m = max(n,256)이어야하는건 이전에 설명한대로 rank 수가 문자 수만큼 나오는게 아니기 때문이라는걸 이해했을거고

 

counting[0] = d로 시작하는게 매우 중요하다.

 

이는 i번째 접미사와 i+d번째 접미사의 rank차이가 최소 d이기 때문이다.

 

 

정렬을 수행하는 순간에, 이전 단계에서 d/2개의 문자를 기준으로 정렬된 상태가 들어오기 때문에..

 

i번째 접미사와 i+d번째 접미사의 글자 차이가 d개니까 최소 d개가 서로 다르다

 

누적합 배열을 만들면 두번째 원소 기준으로 정렬된 인덱스를 저장할 second배열을 만들어야한다

 

rank의 뒤인 n-1부터 순회해서.. i+d번째 원소가 두번째 원소니까 rank[i+d]로 접근할건데..

 

i+d가 n이상이면 빈문자열로 rank = 0으로 주기로 했으므로 rank[n]에 접근한다.

 

초기화할때 rank[n] = 0으로 초기화했고 지금까지 이 부분은 건드리지 않았다..

 

i+d가 n을 넘지 않으면 rank[i+d]에 접근하면 된다.

 

이런식으로 counting sort를 수행하면 된다

 

누적합 배열의 값을 먼저 1 감소하고 (혹은 second[~~-1] = i하고 1 감소시키기) 

 

second[누적합 배열의 값] = 현재 순회중인 index i로 저장

 

second = [0]*n

for i in range(n-1,-1,-1):

    if i+d >= n: #index가 넘어가는 경우, 추가로 만든 n번 index로 

        ind = n

    else:

        ind = i+d

    counting[rank[ind]] -= 1
    second[counting[rank[ind]]] = i

 

 

다음 첫번째 원소를 기준으로 counting sort를 수행

 

i번째 접미사의 rank가 첫번째 원소니까 0부터 n-1까지 rank를 순회해서 counting 배열을 구성하고

 

누적합 배열을 구성한다

 

#첫번째 원소에 대한 정렬

counting = [0]*(m+1)

for i in range(n):

    counting[rank[i]] += 1

for i in range(1,m+1):

    counting[i] += counting[i-1]

for i in range(n-1,-1,-1):

    counting[rank[second[i]]] -= 1

    suffix_array[counting[rank[second[i]]]] = second[i]

return suffix_array

 

그 다음 두번째 원소를 기준으로 정렬된 인덱스 배열인 second 배열을 이용해서 최종 정렬을 시키는데..

 

https://deepdata.tistory.com/968

파이썬 pair type stable sort 배우기(tuple counting sort)

 

second배열의 마지막 인덱스 n-1부터 순회해서.. second[i]값을 구한다.

 

두번째 원소의 인덱스랑 첫번째 원소의 인덱스는 같잖아

 

[첫번째 원소,두번째 원소]를 정렬하는건데 second배열에는 두번째 원소를 기준으로 정렬된 배열의 index가 들어있거든

 

아무튼 정렬하고자 하는 배열인 rank배열의 rank[second[i]]에 접근하고..

 

이는 위 글에서 나타난 "첫번째 원소의 값"이다. 

 

이것을 인덱스로한 누적합 배열에 접근 counting[rank[second[i]]]

 

얘를 먼저 1 감소시키고, counting[rank[second[i]]] -= 1

 

(혹은 counting[rank[second[i]]]-1번 suffix_array에 second[i]를 저장하고, counting[rank[second[i]]]-=1을 하면 되는데 똑같은거잖아)

 

새롭게 정렬할 배열은 바로 suffix_array이다

 

여기서 굳이 새로운 배열 new_rank를 만들 필요가 없다는 것을 이해하는게 중요하다.

 

rank배열을 정렬한 인덱스들의 모임이 suffix_array거든

 

아무튼 suffix_array[counting[rank[second[i]]]]에 정렬 전 배열의 인덱스인 second[i]를 저장하면 된다

 

(저 글에서는 원소를 저장했지만.. 여기서 필요한건 인덱스거든)

 

 

 

2-4) rank를 새롭게 재구성

 

counting_sort를 완성했으니, 이제 d = 1부터 d < 2*n일때까지

 

counting_sort를 하고

 

rank를 새롭게 재구성하고

 

d를 2배하는 것을 반복하는데

 

def get_suffix_array(suffix_array,rank):
    
    d = 1

    while d < 2*n:
        
        suffix_array = counting_sort(suffix_array,rank,d)

        new_rank = [0]*(n+1)

        new_rank[suffix_array[0]] = 1

        for i in range(1,n):
            
            new_rank[suffix_array[i]] = new_rank[suffix_array[i-1]] + compare(rank,suffix_array[i-1],suffix_array[i],d)
        
        #모든 rank가 서로 다르다면 정렬이 끝났다
        if new_rank[suffix_array[n-1]] == n:
            
            break

        rank = new_rank

        d *= 2
    
    return suffix_array

 

 

rank를 재구성하는 부분을 간단하게 할 수 있다

 

"rank배열의 정렬된 결과인 suffix_array가 사실 rank배열과 동일하다"

 

"suffix_array에는 접미사의 정렬된 index가 0~n-1 순서대로 저장되어 있다"

 

그러니까 new_rank = [0]*n으로 초기화하고 가장 빠른 위치인 suffix_array[0]의 rank는 1로 부여하는 것이다.

 

new_rank = [0]*(n+1)

new_rank[suffix_array[0]] = 1

for i in range(1,n):

    new_rank[suffix_array[i]] = new_rank[suffix_array[i-1]] + compare(rank,suffix_array[i-1],suffix_array[i],d)

 

그 다음 i = 1~n-1까지 suffix_array[i]의 rank를 suffix_array[i-1]와 비교해서 새롭게 만들건데..

 

suffix_array[i-1]은 이미 계산되어있다.

 

suffix_array[i-1]과 suffix_array[i]의 rank를 비교해서 서로 같다면 suffix_array[i]는 i-1번과 동일한 rank

 

서로 다르다면 suffix_array[i]는 i-1번 rank에 1을 더한 값을 부여했다.

 

이를 이용해 compare함수를 정의하면..

 

def compare(rank,a,b,d):
    
    if rank[a] == rank[b] and rank[a+d] == rank[b+d]:

        return 0

    else:

        return 1

 

두 접미사 suffix_array[i]와 suffix_array[i-1]의 rank를 비교할때는 첫번째 원소와 두번째 원소가 서로 같으면 동일한 rank를 부여했잖아

 

마찬가지다.. rank배열에는 첫번째 원소는 rank[i]에 있고 두번째 원소는 rank[i+d]에 있다.

 

그러니까 rank[suffix_array[i]]와 rank[suffix_array[i-1]]이 서로 같으면 첫번째 원소가 서로 같은거고

 

rank[suffix_array[i]+d]와 rank[suffix_array[i-1]+d]가 서로 같으면 두번째 원소가 서로 같은 것이다.

 

 

2-5) 근데 index에러가 나야할것 같은데 안나네..?

 

a+d랑 b+d가 n-1을 넘어갈수도 있는거 아니야??

 

넘어갈수도 있는게 맞는데..

 

다음과 같이 rank[a+d]와 rank[b+d]를 먼저 비교하면 index error가 나더라고..

 

근데 rank[a]와 rank[b]를 먼저 비교하면 index error가 나는 경우는 없는듯..?(아무리 해봐도 못찾겠는..?)

 

흐음..?

 

def compare(rank,a,b,d):
    
    if rank[a+d] == rank[b+d] and rank[a] == rank[b]:

        return 0

    else:

        return 1

 

아무튼 첫번째 원소와 두번째 원소의 비교로 new_rank를 구성하고...

 

그런데 new_rank에서 마지막으로 구성된 원소인 new_rank[suffix_array[n-1]]이 n이 된다면...

 

(여기서 new_rank[n-1]이랑 new_rank[suffix_array[n-1]]은 당연히 다르다..)

 

모든 rank가 서로 다르다는 뜻으로 정렬이 완전히 된 상태이니 break하면 된다

 

아니라면 rank = new_rank로 넘겨주고, d를 2배한다음 다음 반복으로..

 

반복이 끝나면 suffix_array에 정렬된 결과가 있으니 이를 return하면 된다

 

 

2-5) kasai 알고리즘..

 

kasai 알고리즘은 잘 구현된 것 같다 

 

전체 속도도 빠른거보면 여기는 문제는 없는듯

 

 

2-6) 빠른 출력

 

배열을 출력할때 단순히 print(*arr)보다 문자열로 바꿔서 join하면 더 빠른 출력이 가능하다고 한다

 

실제로 0.1초 정도 더 빠르긴함

 

#배열에 대한 빠른 출력
print(" ".join(map(lambda i:str(i+1), suffix_array)))
print(" ".join(map(str,lcp)))

 

https://gumgood.github.io/python-fastio

Python FastIO - 프로그래밍 대회를 위한 파이썬 빠른 입출력 - Gumgood

 

 

3. $O(nlog^{2}n)$ 최적화 코드

 

counting sort 부분을 sorted()함수로 퀵정렬하면 $O(nlog^{2}n)$최적화된 코드를 얻는다

 

counting sort를 안해도 되니까 구현이 조금 더 간단

 

대신 조금 더 느리다(n = 10만 기준 counting sort는 0.2초 quick sort는 1초정도)

 

def compare(rank,a,b,d):
    
    if rank[a] == rank[b] and rank[a+d] == rank[b+d]:
        
        return 0
    
    else:
        
        return 1

def get_suffix_array(rank):
    
    d = 1

    while d < 2*n:

        suffix_array = sorted(range(n), key = lambda i: [rank[i],rank[min(i+d,n)]])

        new_rank = [0]*(n+1)

        new_rank[suffix_array[0]] = 1

        for i in range(1,n):
            
            new_rank[suffix_array[i]] = new_rank[suffix_array[i-1]] + compare(rank, suffix_array[i-1], suffix_array[i] , d)

        if new_rank[suffix_array[n-1]] == n:
            
            break
        
        rank = new_rank

        d *= 2

    return suffix_array 
            

s = input()

n = len(s)

rank = [ord(s[i]) - ord('a') + 1 for i in range(n)]
rank.append(0)

suffix_array = get_suffix_array(rank)

for i in suffix_array:
    
    print(i)

 

suffix array를 초기화하지 않아도 되고,

 

counting sort하는 부분을 다음과 같이 고친다.

 

suffix_array = sorted(range(n), key = lambda i: [rank[i],rank[min(i+d,n)]])

 

0~n-1을 정렬하는데, rank[i] 값에 따라 정렬을 한다. rank[i]가 같으면 rank[min(i+d,n)]에 따라 정렬한다.

 

rank[i]는 i번째 접미사의 rank로 rank배열의 첫번째 원소

 

rank[min(i+d,n)]은 i+d번째 접미사의 rank이며 i+d가 n 이상이면 rank[n]으로 0을 부여해준다.

 

suffix_array는 rank배열의 정렬된 인덱스이기때문에, 인덱스를 얻기 위해 위와 같이 정렬한다는 것을 기억

 

 

4. 연습문제

 

13264번: 접미사 배열 2 (acmicpc.net)

13264번: 접미사 배열 2

 

n이 최대 10만에서 접미사 배열만을 구하는 문제