suffix array 깊이 이해하기 2 -부분 문자열의 개수를 세는 방법-

1. 문제1

 

11479번: 서로 다른 부분 문자열의 개수 2 (acmicpc.net)

11479번: 서로 다른 부분 문자열의 개수 2

 

2. 풀이

 

1) "모든 부분 문자열은 접미사들의 접두사이다"

 

suffix array를 구한다면 suffix_array[i]는 사전순으로 오름차순 정렬된 i번째 접미사의 위치를 알 수 있다

 

그러면 문자열의 길이가 n이라면 i번째 접미사의 길이는 얼마일까? 당연히 n - suffix_array[i]이다.

 

"ABAAB"에 대하여 suffix array를 생각해보자.

 

index	suffix array
2	AAB
3	AB
0	ABAAB
4	B
1	BAAB

 

 

오름차순으로 정렬된 1번째 접미사는 suffix_array[1] = 3이므로...

 

각 접미사의 길이가 n - suffix_array[i]이므로, 각 접미사의 모든 접두사는 n-suffix_array[i]개가 존재한다.

 

예를 들어 길이 5인 ABAAB는 접두사가 A, AB, ABA, ABAA, ABAAB로 5개이다.

 

즉 문자열은 길이만큼 서로 다른 접두사를 가진다.

 

모든 부분문자열은 접미사의 접두사이므로 이러한 접두사의 개수를 세면 부분문자열의 개수를 알 수 있게된다.

 

그런데 원하는건 서로 다른 부분 문자열의 개수이다.

 

당연히 동일한 부분 문자열이 있을 것인데... 여기서 매우 중요한 성질 하나를 익히고 가면..

 

 

2) i번째 접미사는 i-1번째 접미사와 lcp[i]만큼 서로 겹친다.

 

당연한 소리지만.. i번째 접미사는 이웃한 이전 접미사 i-1번째 접미사와 lcp[i]만큼의 공통 접두사를 가진다.

 

접미사의 접두사는 부분문자열이므로 lcp[i]만큼 서로 겹치는 부분문자열이 존재한다는 뜻이다.

 

실제로 위의 경우를 구해보자.

 

index	suffix array	lcp
2	AAB	-
3	AB	1
0	ABAAB	2
4	B	0
1	BAAB	1

 

그러면 AAB부터 부분 문자열의 개수를 세보면..?

 

1) AAB는 A, AA, AAB로 3개

 

 

2) 다음으로 넘어가서 AB는 A, AB로 2개인데.. lcp값이 1이다. 

 

즉 앞에서 센 부분문자열중 1개가 동일한 부분문자열이 있다는 뜻이다. 

 

실제로 A가 서로 겹친다. 그러므로 3 + 2 - 1 = 4개

 

 

3) 다음으로 넘어가면 ABAAB로 A,AB,ABA,ABAA,ABAAB로 5개 있는데 lcp값이 2이다.

 

이는 이전 접미사 AB와 겹치는 부분문자열이 2개 있다는 뜻이다. 

 

실제로 A,AB가 서로 겹치고 있다. 그러므로 A,AA,AAB,AB,ABA,ABAA,ABAAB로 3 + 2 - 1 + 5 - 2 = 7개

 

 

4) 다음으로 넘어가서 B는 부분문자열이 1개 있는데 lcp가 0으로 이전에 센 부분문자열과 겹치는게 없다.

 

실제로도 ABAAB의 부분문자열과 겹치는 문자열이 존재하지 않는다.

 

그러므로 A,AA,AAB,AB,ABA,ABAA,ABAAB,B로 8개

 

 

5) 다음 BAAB는 B,BA,BAA,BAAB로 4개있는데 lcp값이 1로 이전 B와 겹치는 부분문자열이 1개 있다는 뜻이다.

 

실제로 B가 서로 겹치고 있다. 

 

그러므로 A,AA,AAB,AB,ABA,ABAA,ABAAB,B,BA,BAA,BAAB로 11개

 

----------------------------------------------------------------------------------------------------------------------------------------

 

그러므로 서로 다른 부분 문자열의 개수를 세는 구체적인 알고리즘은..

 

suffix array와 lcp배열을 구한다.

 

suffix array를 순회하면서 접미사의 길이 n-suffix_array[i]를 더해간다.

 

1번 원소부터는 lcp[i]만큼 겹치므로 n-suffix_array[i] - lcp[i]를 더해나간다.

 

즉, n - sa[0] + (n-sa[1]-lcp[1] + n-sa[2]-lcp[2] + n-sa[3]-lcp[3]+...+n-sa[n-1]-lcp[n-1])

 

여기서 suffix_array의 원소 합은 0부터 n-1까지의 합이므로 (n-1)n//2이다.

 

$$n*n - \sum_{i=0}^{n-1}sa[i] - \sum_{i=1}^{n-1}lcp[i] = n*n - \frac{(n-1)n}{2} - \sum_{i=1}^{n-1}lcp[i]$$

 

식을 정리하면...

 

$$\frac{n^{2}+n}{2} - \sum_{i=1}^{n-1}lcp[i]$$

 

#서로 다른 부분문자열의 개수 구하기

#suffix array를 구하기
def counting_sort(suffix_array,rank,d):
    
    m = max(n,256)

    counting = [0]*(m+1)

    counting[0] = d

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

    for i in range(n-1,-1,-1):
        
        if i+d >= n:
            
            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(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)
        
        if new_rank[suffix_array[n-1]] == n:
            
            break
        
        rank = new_rank
        d *= 2
    
    return suffix_array

#lcp 배열을 구하기
def kasai(suffix_array,s):
    
    lcp = [0]*n
    reverse = [0]*n

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

    for i in range(n):
        
        if reverse[i] == 0:
            
            continue
        
        adj = suffix_array[reverse[i]-1]

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

        if k > 0:
            
            k -= 1
    
    return lcp

s = input()

n = len(s)

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

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

#모든 부분문자열은 접미사의 접두사이다.
#정렬된 i번째 접미사의 길이는 n - suffix_array[i]
#각 접미사는 n - suffix_array[i]개만큼 접두사(=부분문자열)를 가진다.
#i번째 접미사는 i-1번째 접미사와 lcp[i]만큼 서로 동일하다.
#따라서 0~n-1번 접미사를 순회하면서 부분문자열 개수를 더해나가고 겹치는 개수만큼 뺀다면..
#n - sa[0] + n-sa[1]-lcp[1] + n-sa[2]-lcp[2] + ... + n-sa[n-1]-lcp[n-1]
#n(n+1)//2 - sum(lcp)

print(n*(n+1)//2 - sum(lcp))