백준1787 문자열의 주기 예측

작성일 | In

문제 링크

  • http://icpc.me/1787

문제 출처

  • 2005/2006 POI Stage1 2번

사용 알고리즘

  • KMP

시간복잡도

  • $O(N)$

풀이

0 ≤ i ≤ N-1인 i에 대해, S[0..i]에서 prefix와 suffix가 같은 가장 짧은 길이를 구하는 문제입니다.

일단, KMP의 failure function을 잘 생각해보면, 재귀적이라는 사실을 알 수 있습니다.
예를 들어, N = 8이고 pi[N-1] = 4이면 S[0..3]과 S[4..7]이 같습니다.
만약 이때 pi[3] = 3이라면, S[0..2]와 S[1..3]이 똑같고, S[4..6]과 S[5..7]도 같습니다.
그러므로 S[0..2]와 S[5..7]이 같다는 것을 알 수 있고, 이렇게 재귀적으로 타고 내려가면 prefix = suffix인 문자열들을 모두 볼 수 있습니다.

메모이제이션을 해주면 O(N)에 풀 수 있습니다.

전체 코드

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
#include <bits/stdc++.h>
using namespace std;

typedef long long ll;

int n; string s;
int pi[1010101];
ll dp[1010101];

ll f(int x){
    if(x < 0) return 1e18;
    ll &res = dp[x];
    if(res != -1) return res;
    if(!pi[x]) return res = 1e18;
    return res = min<ll>(pi[x], f(pi[x]-1));
}

int main(){
    ios_base::sync_with_stdio(0); cin.tie(0);
    cin >> n >> s;
    for(int i=1,j=0; i<n; i++){
        while(j && s[i] ^ s[j]) j = pi[j-1];
        if(s[i] == s[j]) j++;
        pi[i] = j;
    }
    ll sum = 0;
    memset(dp, -1, sizeof dp);

    for(int i=0; i<n; i++){
        ll t = f(i);
        //cout << i << " : " << t << "\n";
        if(t < 1e18) sum += i-t+1;
    }
    cout << sum;
}