백준 11479 서로 다른 부분 문자열의 개수 2
23 Jun 2021https://www.acmicpc.net/problem/11479
문자열 S의 부분 문자열 중에서 서로 다른 것의 갯수를 구하는 문제이다. 예시를 들어서 이해해보자. 먼저 ababc라는 문자열이 주어진다고 했을 때 5개의 접두사를 얻어낼 수 있다.
- a
- ab
- aba
- abab
- ababc
그런데 맨 앞에 a를 탈락시킨 babc에 대해서도 접두사를 구하면 4개의 접두사를 얻을 수 있다.
- b
- ba
- bab
- babc
이런식으로 접미사들의 모든 접두사를 구하면 우리가 구하고자 하는 부분 문자열을 모두 나타낼 수 있다. 그런데 이중에서 겹치지않는 것을 찾고자하려면 어떻게 해야할까? 생각해보면 접미사들의 접두사를 구하는 것으로 부분 문자열을 나타낸다고 했었다. 만약 두 접미사의 앞부분이 일치한다면 접두사를 구하는 과정에서 부분 문자열이 중복해서 나타날 것이다. 예시로 ababc를 다시 보자.
| Suffix Array | i | LCP |
|---|---|---|
| ababc | 0 | x |
| abc | 2 | 2 |
| babc | 1 | 0 |
| bc | 3 | 1 |
| c | 4 | 0 |
ababc에 대한 접두사들을 모두 구하고 다음 suffix에 대해 새로운 접두사를 구하려고 한다고 하자. abc가 ababc와 중복되지않는 부분문자열의 갯수는 abc의 길이 - LCP[abc] 로 3-2가 되어서 1이 나온다. 실제로도 a와 ab는 이미 등장했기 때문에 abc만 세주면 된다.
만약 abc라는 부분 문자열이 다시 등장하려면 abcd와 같이 abc 뒤에 무언가 이어진 형태가 나타나야할 것이고 이는 suffix상에서 인접해있기 때문에 lcp를 구하는 과정에서 포함이 된다. 따라서 abc는 두번 카운트 되지 않는다. 또한 dabc와 같은 형태는 무조건 접두사를 구하는 과정에서 앞에 d를 포함시키게 된다. 따라서 abc는 두번 카운트 되지 않는다.
최종적으로 LCP와 Suffix Array를 구해서 다음의 작업을 해주면 서로 다른 부분 문자열의 개수를 구할 수 있다. 이 때 LCP[0]의 값을 0으로 설정해주도록 하자. n-sa[i]를 하면 i번째 suffix의 길이가 나온다는 것을 명심하자.
for i in 0..<n {
result += n-sa[i]-lcp[i]
}
- 전체 코드
func compare(g: [Int], a: Int, b: Int, t: Int) -> Bool {
if g[a] == g[b] {
return g[a+t] < g[b+t]
}
else {
return g[a] < g[b]
}
}
func suffix(str: [Character]) -> [Int] {
let n = str.count
var g = Array(repeating: -1, count: n+1)
var ng = Array(repeating: -1, count: n+1)
var sfx = Array(repeating: -1, count: n)
var t = 1
for i in 0..<n {
g[i] = Int(str[i].asciiValue!-96)
sfx[i] = i
}
while t < n {
sfx.sort(by: {compare(g: g, a: $0, b: $1, t: t)})
ng[sfx[0]] = 0
for i in 1..<n {
if compare(g: g, a: sfx[i-1], b: sfx[i], t: t) {
ng[sfx[i]] = ng[sfx[i-1]] + 1
}
else {
ng[sfx[i]] = ng[sfx[i-1]]
}
}
g = ng
t <<= 1
}
return sfx
}
func LCP(str: [Character], SA: [Int]) -> [Int] {
let n = str.count
var strArr = str
strArr.append("#")
var lcp = Array(repeating: 0, count: n)
var rank = Array(repeating: 0, count: n)
for i in 0..<n {
rank[SA[i]] = i
}
var len = 0
for i in 0..<n {
let k = rank[i]
if k > 0 {
while strArr[SA[k-1]+len] == strArr[i+len] {
len += 1
}
lcp[k] = len
if len > 0 {
len -= 1
}
}
}
return lcp
}
let line = readLine()!.map{Character(String($0))}
let sa = suffix(str: line)
let lcp = LCP(str: line, SA: sa)
let n = line.count
var result:UInt64 = 0
for i in 0..<n {
result += UInt64(n-sa[i]-lcp[i])
}
print(result)