1787 - 문자열의 주기 예측

문제 링크

---

https://www.acmicpc.net/problem/1787

문제 요약

---

문제가 이해를 하기 조금 힘들지만 결국 하는말은 다음과 같다. 어떠한 문자열 \(X\)가 있다고 가정을 하자. 이 때 \(X\)의 \(Prefix\)와 \(Suffix\)를 구하는데 이 때 빈 문자열이면 안되고 서로 동일한 길이여야 한다. 이렇게 뽑았을 때 \(Prefix\)와 \(Suffix\)가 같은 0이 아닌 가장 작은 길이를 구한다음에 \(X\)의 길이에서 뺀뒤 이 값을 답에 누적한다. 이러한 과정을 문자열 \(S\)가 주어지면 \(S\)의 모든 \(Prefix\)에 대해서 반복한다.

문제 풀이

---

주어지는 문자열 \(S\)의 길이가 \(10^6\)이기 때문에 빠른 풀이를 생각해야 한다. 결론부터 얘기하면 이 문제는 \(KMP\)를 이해하고 있으면 쉽게 풀 수 있다. 일단 위에서 말 한 임의의 문자열 \(X\)에 대해서 \(Prefix\)와 \(Suffix\)가 동일하게 만드는 0이 아닌 가장 작은 길이를 어떻게 구하는지 생각을 해 보자. 만약 이 문제를 해결할 수 있다면 문자열 \(S\)의 모든 \(Prefix\)에 대해서 문제를 확장해 우리가 원하는 답을 구할 수 있을 것이다.

우선 \(failure\ function\)이 어떠한 의미를 가지는지 생각을 하자. \(pi[|X| - 1]\)에 만약 5라는 값이 써져 있다면 문자열 \(X\)는 \(Prefix\)와 \(Suffix\)가 최대 5만큼 똑같다는 말이 된다. 아쉽게도 이값은 우리가 구하려는 값이 아니다. 우리가 구하려는 값은 최대로 겹치는 수가 아니라  최소로 겹치는 수를 구해야 한다. 그렇다면 이 값을 어떻게 이용해야 할까.

\(Prefix\)와 \(Suffix\)가 동일하다는 점에 주목해보자. \(X\)의 길이가 7이고 \(Pi\)배열의 마지막 수가 5라면 \(X_1...X_5\)는 \(X_3...X_7\)과 같다. 만약 \(X_5\)의 \(Pi\)배열 값이 2라면 어떻게 될까? \(X_1 = X_4,\ X_2 = X_5\)가 된다. 헌데 \(X_4 = X_6,\ X_5 = X_7\)이다. 즉 \(X_1 =  X_6,\ X_2 = X_7\)이 성립한다. 이와같은 과정을 반복해 \(X\)에서 \(Prefix\)와 \(Suffix\)가 똑같은 0이 아닌 최소의 길이를 구할 수 있다.

이것을 \(S\)의 모든 \(Prefix\)로 확장하는 것은 그리 어렵지 않다. \(S\)의 \(Pi\)배열을 미리 구해놓고 차례대로 위와같은 과정을 진행하면 된다. 하지만 매번 계산할 때 겹치는 계산구간이 많으므로 메모이제이션을 하면 \(O(N)\)의 시간으로 문제를 해결할 수 있다.

소스 코드

---

#include <stdio.h>
#include <string.h>
#include <algorithm>

using namespace std;

typedef long long ll;

const int INF = 987654321;
const int MAX = 1e6 + 10;
int n, dp[MAX], pi[MAX];
char str[MAX];

int func(int pos) {
  if (pos < 0 || pi[pos] == 0) {
    return INF;
  }
  
  int& ret = dp[pos];
  if (ret != - 1) {
    return ret;
  }
  
  return ret = min(pi[pos], func(pi[pos] - 1));
}

int main() {
  scanf("%d%s", &n, str);
  
  int match = 0;
  for (int i = 1 ; i < n ; ++i) {
    while (match && str[i] != str[match]) {
      match = pi[match - 1];
    }
    if (str[i] == str[match]) {
      ++match;
    }
    pi[i] = match;
  }
  
  memset(dp, -1, sizeof(dp));
  ll ans = 0;
  for (int i = 0 ; i < n ; ++i) {
    int ret = func(i);
    if (ret != INF) {
      ans += (ll)(i - ret + 1);
    }
  }
  
  printf("%lld\n", ans);
}