KMP 알고리즘으로 부분 문자열 효과적으로 제거하기

이 포스트에서는 USACO 문제인 Censoring에서 문제 풀이를 통해 KMP와 동적 계획법으로 특정한 부분 문자열을 전체 시간 복잡도 $O(N + M)$로 처리하는 방법을 알아본다.

 

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

10747번: Censoring

목표

어떤 문자열이 주어지고 삭제하고자 하는 문자열이 주어지면 문자열에서 해당 부분 문자열이 나타나지 않을 때까지 그 부분 문자열을 모두 삭제해야 한다.

 

예를 들어 whatthemomooofun 에서 나타나는 모든 moo를 삭제하고자 하면 그 결과는 whatthefun이 된다.

단순한 방법

이를 해결하는 단순한 방법은 스택 구조를 사용해 $O(NM)$의 시간 복잡도로 해결하는 방법이다.

 

현재 인덱스 i에 대해 s[i]가 패턴의 맨 뒤와 동일하면 패턴의 길이 M만큼 거슬러 올라가 완벽히 같으면 M개의 문자를 pop 하는 방식으로 해결할 수 있다.

 

이는 COCI 문제인 문자열 폭발에서 푸는 방법이 널리 알려진 상태이다. 하지만 위에 먼저 언급한 Censoring의 경우 N과 M 둘 다 $10^{6}$이기 때문에 이 방법으로는 풀 수 없음이 너무나도 당연하다.

 

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

9935번: 문자열 폭발

동적 계획법이요?

사실 우리는 KMP 알고리즘에서 실패 함수를 구축하는 것 자체가 메모이제이션의 성질을 띄고 있기 때문에 무슨 동적 계획법을 또 하느냐는 의문을 가질 수 있다.

 

우리는 테이블에서 다음 정보를 저장할 것이다.

 

해당 인덱스를 매칭 했을 때 계산되는 공통 매칭 값(소위 j로 두는 그 값이다.)

 

이를 저장하고 다닌다고 했을 때 위 예시를 풀어보자.

 

whatthemo까지 했을 때 테이블이 저런 식으로 채워지리란 건 쉽게 유추할 수 있다. 그 뒤에 있는 ...moo까지 매칭 해보자.

 

 

이제 여기서 j = 3이 되고 moo가 매칭 되었음을 확인했다. 이제 우리는 여기서 moo를 제거할 것이다.

 

 

삭제한 건 좋은데 이다음 글자가 o가 되고 또 매칭이 되리란 건 눈으로 봐도 알 수 있지만 그 사실을 어떻게 알려줘야 할지 잘 모르는 상황이다. 여기서 테이블에 저장해둔 값들이 진가를 발휘한다.

 

 

moo를 삭제했을 때 가장 마지막 문자가 o고 해당 인덱스의 테이블 값은 2이다. 저장된 테이블 값을 j에 넘겨줌으로써 "그래서 얘는 어디까지 매칭 되었더라?"를 알려주게 된다. 그러면 j의 값은 2가 되고 바로 다음 글자 o가 매칭 되어 j + 1 = 3이므로 또 다른 moo 발견, 이렇게 모든 moo를 찾게 되는 것이다.

구현

구현은 크게 어렵지 않고 결과로 따로 만들 문자열과 dp 테이블을 추가하고 KMP에 약간의 수정을 가하면 된다.

string text, pat;
string ans;
int dp[1000000];
 
vector<int> bake_pi(string &pt)
{
  vector<int> pi(pt.size());
  int j = 0;
  for (int i = 1; i < pt.size(); i++)
  {
    while (j and pt[j] != pt[i]) j = pi[j - 1];
    if (pt[i] == pt[j]) pi[i] = ++j;
  }
  return pi;
}
 
void kmp(string &target, string &pt)
{
  auto pi = bake_pi(pt);
  int j = 0;
  for (int i = 0; i < target.size(); i++)
  {
    ans.push_back(target[i]);
    while (j and pt[j] != ans.back()) j = pi[j - 1];
    if (pt[j] == ans.back()) ++j;
    if (j == pt.size())
    {
      for (int k = 0; k < pt.size(); k++)
      {
        dp[ans.size() - 1] = 0;
        ans.pop_back();
      }
      j = 0;
      if (ans.size())
      {
        j = dp[ans.size() - 1];
      }
    }
    else
    {
      dp[ans.size() - 1] = j;
    }
  }
}

 

코드를 보면 ans와 dp가 새로 추가되었다. KMP를 수행하면서 원래 문자열인 target에 대해 비교하는 게 아닌 ans에 일단 하나 넣고 ans.back()과 비교하여 j를 계산하는 것을 볼 수 있다.

 

그러다 문자열 매칭이 성공하면 패턴의 길이만큼 ans에서 pop을 하면서 해당하는 인덱스의 테이블 값을 0으로 다시 비워주게 된다.

 

그리고 j를 기존에 저장했던 값으로 바꿔치기를 하고 계속 탐색한다. 물론 매칭이 끝나지 않으면 해당 j 값을 dp 배열에 저장하여 훗날 값을 꺼내 쓸 수 있도록 해준다.

 

KMP의 수행이 끝나면 ans에 우리가 원하는 결과가 들어있게 된다.

 

이렇게 우리는 특정 패턴을 삭제하는 문제를 KMP로 아주 빠르게 푸는 방법을 배워봤다.

 

실제로 활용할 수 있는 문제가 나온 만큼 방법을 잘 숙지하면 스택으로 풀 수 없는 제한이 나왔을 때 여러분에게 큰 무기가 되어줄 수 있을 것이다.