BOJ 13576. Prefix와 Suffix

13576번: Prefix와 Suffix (acmicpc.net)

solved.ac: Platinum II (2024.09.10)

 

꽤나 흥미롭게 풀었던 문제입니다.

배울 점이 많은 문제라고 생각합니다.

 

1. Prefix이자 Suffix인 substr 구하기

처음엔 전체 문자열 하나 잡고, 앞으로 계속 lcp랑 len이랑 다를 때까지 탐색하면 되지 않을 까 생각했습니다.

그러나 BBAABB같은 문자열을 생각해 보면, 씨알도 안 먹히는 소리임을 알 수 있습니다.

 

처음엔 쌩 SA + LCP만 가지고 풀어보려 했으나, 능지 이슈로 이건 아직도 모르겠습니다.

그래서 그냥 실패함수 가져다 썼습니다. 조건을 만족시키는 substr의 길이는 fail(len)을 반복적으로 수행하여 얻을 수 있습니다.

 

2. 그래서 문자열 나타는 거 어케 확인함?

어떤 문자열에 대해서, 문자열이 나타나는 지의 여부를 판단하는 방법을 생각해 봅시다.

문자열 S에 대해 S[i]를 i번째 글자부터 시작하는 접미사라 생각해 보면,

모든 접미사에 대해 해당 문자열이 접두사가 되는 지 확인하면 됩니다.

 

물론 나이브하게 구현하면 당연히 노답임을 알 수 있습니다. 우리가 접두사인지 판단해 보는 문자열도 접미사임을 상기합시다.

확인하고자 하는 접미사 T와 다른 접미사에 대해 공통 접두사의 길이가 len(T)라면, 그 접미사는 T를 접두사로 갖습니다.

 

이때 T를 접두사로 갖는 접미사 A와 다른 접미사 B에 대해, 공통 접두사의 길이가 len(T) 이상이라면,

B 역시 T를 접두사로 갖는 접미사임을 알 수 있게 됩니다.

 

여기까지 왔다면 뭘 써야 하는지는 지레짐작이 갑니다. 접미사 배열을 만들어야 합니다.

접미사 배열은 정렬되어 있으므로, 인접한 두 접미사끼리 LCP를 비교하면서 한 번이라도 길이가 len(T) 미만으로 간다?

그 이후의 문자열은 볼 필요도 없습니다. 딱 거기까지만 개수를 세면 됩니다.

 

map 등을 이용해서 길이를 통해 인덱스를 O(logN) 미만으로 구할 수 있게 구성하면,

N이 10만으로 상당히 애매하기 때문에 O(N²)이 어떻게 잘 통과됩니다. 2초라서 7억 번 정도는 잘 돌아가는 것 같습니다.

 

아래는 첫 풀이입니다. (1540ms)

#include <iostream>
#include <vector>
#include <algorithm>
#include <string>
#include <cstring>
#include <map>

using namespace std;

/* returns pair of suffix array, lcp array. lcp range = [0, n - 1] */
pair<vector<int>, vector<int>> sa_lcp(string str)
{
  int n = str.length(), d = 1;
  vector<int> pos(n), sa(n), lcp(n - 1), tmp(n);

  auto cmp = [&n, &d, &pos](int i, int j)
  {
    if(pos[i] != pos[j]) return pos[i] < pos[j];
    i += d; j += d;
    return (i < n && j < n) ? (pos[i] < pos[j]) : (i > j);
  };

  //struct suffix array
  for(int i = 0; i < n; i++) sa[i] = i, pos[i] = str[i];
  for(d = 1; ; d <<= 1)
    {
      sort(sa.begin(), sa.end(), cmp);
      for(auto &i: tmp) i = 0;
      for(int i = 0; i < n - 1; i++)
        {
          tmp[i + 1] = tmp[i];
          if(cmp(sa[i], sa[i + 1]) == true) tmp[i + 1]++;
        }
      for(int i = 0; i < n; i++) pos[sa[i]] = tmp[i];

      //모든 접미사가 서로 다른 그룹에 존재
      if(tmp[n - 1] == n - 1) break;
    }

  //struct LCP array
  for(int i = 0, k = 0; i < n; i++, k = max(k - 1, 0))
    {
      //마지막 접미사, 비교할 접미사 존재하지 않음
      if(pos[i] == n - 1) continue;
      //다음 그룹의 접미사
      for(int j = sa[pos[i] + 1]; str[i + k] == str[j + k]; k++);
      lcp[pos[i]] = k;
    }

  return {sa, lcp};
}

vector<int> failureFunction(string substr)
{
  int sublen = substr.length();
  vector<int> pi(sublen);
  int prefix = 0;
  for(int suffix = 1; suffix < sublen; suffix++)
    {
      while(prefix > 0 && substr[prefix] != substr[suffix]) prefix = pi[prefix - 1];
      if(substr[prefix] == substr[suffix]) pi[suffix] = ++prefix;
    }
  return pi;
}

/* len -> i */
map<int, int> idx;

int main()
{
  ios::sync_with_stdio(false);
  cin.tie(0);

  string str;
  cin >> str;

  int n = str.length();
  auto arrays = sa_lcp(str);
  auto sa = arrays.first, lcp = arrays.second;
  lcp.push_back(0);

  auto fail = failureFunction(str);
  for(int i = 0; i < n; i++) idx[n - sa[i]] = i;
  
  vector<pair<int, int>> res;
  int len = n;
  while(len)
    {
      int i = idx[len];
      int cnt = 0;
      while(lcp[i + cnt] >= len) cnt++;
      res.push_back({len, cnt + 1});
      len = fail[len - 1];
    }
  
  sort(res.begin(), res.end());

  cout << res.size() << '\n';
  for(auto ans: res)
    {
      cout << ans.first << ' ' << ans.second << '\n';
    }
}

3. 개억지로 푼 거 아님?

일단 풀기는 했지만, 상당히 불-편합니다.

애초에 N 제한부터 O(N²) 풀이를 의도한 것은 아님을 알 수 있습니다.

 

물론 뭐 세그트리라던가 세그트리라던가 세그트리같은 걸 끼워팔아서 O(NlogN)으로 줄일 수는 있겠지만,

이 문제는 약간의 관점 전환으로 훨씬 간단히 해결할 수 있습니다.

 

문제에서 확인할 문자열은 접두사이사 접미사인 부분 문자열들입니다.

여기서 접두사라는 사실에 조금 주목을 해 보면, 길이 L의 부분 문자열에 대해서,

전체 문자열과 각각의 접미사에 대해 LCP가 L 이상이면 해당 접미사의 맨 처음 인덱스에서 매칭이 가능함을 알 수 있습니다.

 

미리 길이들을 다 구해놨다고 하면, 필요한 것은 '각각의 접미사에 대한 전체 문자열과의 LCP 길이'입니다.

이는 Z 알고리즘으로 선형 시간에 구할 수 있습니다.

 

Z array를 활용하면, 사실 실패함수로 길이를 구하지 않아도 됩니다.

문제의 조건에서 접두사이자 접미사인 문자열의 길이를 요구하고 있고, 각각의 접미사에 대해

접두사의 길이와 접미사의 길이가 같다면, 즉 i + z[i] == n이면 구하는 길이는 z[i]가 됩니다.

 

LCP의 길이가 L 이상인 접미사의 개수, 이건 그냥 누적합으로 간단히 구할 수 있습니다.

아래는 2차 코드입니다. 실행 시간은 16ms로 확 줄었습니다.

#include <iostream>
#include <vector>
#include <algorithm>
#include <string>

using namespace std;

/* Z array: 각각의 접미사에 대한 공통 접두사 길이 */
vector<int> zFunction(string &str)
{
  int n = str.length();
  vector<int> z(n);

  z[0] = n;
  int l = 0, r = 0;
  for(int i = 1; i < n; i++)
    {
      if(i < r) z[i] = min(z[i - l], r - i);
      while(i + z[i] < n && str[z[i]] == str[i + z[i]]) z[i]++;
      if(i + z[i] > r) l = i, r = i + z[i];
    }
  return z;
}

int cnt[100001];

int main()
{
  ios::sync_with_stdio(false);
  cin.tie(0);

  string str;
  cin >> str;

  int n = str.length();
  auto z = zFunction(str);

  for(int i = 0; i < n; i++) cnt[z[i]]++;
  for(int i = n - 1; i >= 0; i--) cnt[i] += cnt[i + 1];

  vector<pair<int, int>> res;
  for(int i = n - 1; i >= 0; i--)
    {
      if(i + z[i] == n) res.push_back({n - i, cnt[z[i]]});
    }

  cout << res.size() << '\n';
  for(auto pi: res) cout << pi.first << ' ' << pi.second << '\n';
}