Suffix Arrayを爆速で構築したら、驚きの結果に…？！？！？

Posted on 6 mins read

ちょっとまえにSA-ISを実装したので、メモがてらすごく簡単に記録しておきます。

SA-ISについて

Suffix Arrayを線形時間で構築するアルゴリズムです。かなり賢いアルゴリズムで、2009年に論文が出ています。
suffixを比較するのにO(n)かかるので、単純にソートをするとO(n2 * log(n))になるのはわかると思います。

蟻本に計算量O(n * log2(n))のダブリングという手法を使った実装が載っているのですが、その実装ではこの問題を解いたときにTLEしたので、SA-ISを実装しようと思いました。

アルゴリズムのきもち

もとの文字列Sをいくつかのsubstring($s_1, s_2, …$)に分割します。

ここで、$s_1, s_2, …$の辞書順がわかったとします。(*1)
すると、$s_1, s_2, …$それぞれをまとめて1つのchar $c_1, c_2, …$としてみた場合に、任意の$c_i, c_j$を比較することができます。

さて、$c_1, c_2, …$をそのまま連結してできる文字列をS_newとして、S_newのsuffix arrayを考えてみましょう。(これは再帰で求めます)
そうするとSにおいて、$s_1, s_2, …$の先頭の文字(substringに区切った場所)だけについてですが、suffixの辞書順が得られたことになります！！

$s_1, s_2, …$を適切に区切っていると、ここから全体のsuffix arrayを復元するのは実はそれほど難しくないです。(*2)

1と2の操作を効率的に行うために、文字列のインデックスをS型、L型に分類してinduced sortというバケツソートの一種を実行することがこのアルゴリズムのポイントです。

計算量は、1, 2の操作がともに線形時間でできた場合、O(n) + O(n / 2) + O(n / 4) + … = O(n)となります。 (任意の$s_i$は必ず2文字以上からなるので、再帰が呼ばれるたびに，suffix arrayを求める必要がある文字列の長さは1/2以下になるからです)

アルゴリズム概要

  1. S型、L型、LMSを割り振って、LMSによって文字列を分割する。
    文字列の全てのインデックスについて、そこから始まるsuffixと一つあとから始まるsuffixを比べて増加しているか、減少しているかによってそれぞれS型、L型を割り振ります。 LMS(left most S)は、連続するSにおいてもっとも左にあるSです。
  2. induced sort(1回目)
    特殊なバケツソートを実行します。正しいinduced sortにはLMS-suffixの順序を渡す必要があるのですが、まずは適当な順番で渡して実行してしまいます。これによって得られるものをpseudo suffix arrayとします。
  3. LMS-suffixの順序を求める
    さて、先ほど求めたpseudo suffix arrayについてですが、これはもちろん正しいsuffix arrayではありません。 これが満たす性質を頑張って考えると、実はLMSだけに注目したときに、LMS-substring($s_1, s_2, ...$)についての順序を表していることがわかります。(ここがすごく賢い！頑張って考えるとなんとなくわかります。詳しくは他の記事)
    LMS-substringの順序がわかったので、上に書いたように再帰でLMS-suffixの順序を求めることができます。
  4. induced sort(2回目)
    LMS-suffixの順序がわかったので、これを渡してinduced sortをすると正しいsuffix arrayが求まります。

細かい部分は飛ばしてしまったので下の実装のコメントや他の記事を参考にしてください。
証明は全く書いてないので他の記事参照。これやこれがわかりやすかったです。

実装

verifyはこちら。 可読性を上げることを意識したために冗長な部分が多いです。

#include <bits/stdc++.h>
#define REP(i, n) for (int (i) = 0; (i) < (n); (i)++)
#define REPI(i, a, b) for (int (i) = (a); (i) < (b); (i)++)
using namespace std;
using VI = vector<int>;
class SuffixArray {
  // str: 文字列をintのvectorに直したもの、k: 現れる文字の種類数
  VI sa_is(const VI& str, const int k) {
    const int n = str.size();

    // 1. S, L, LMSを求める
    vector<bool> is_S(n); is_S[n - 1] = true;
    vector<bool> is_LMS(n);
    VI LMSs;
    for (int i = n - 2; i >= 0; i--) {
      // Sである条件は、str[i..] < str[i + 1..] (str[i] == str[i + 1]のときは一つ右に依存)
      is_S[i] = str[i] < str[i + 1] || (str[i] == str[i + 1] && is_S[i + 1]);
    }
    REP (i, n) {
      // LMSは一番左のS
      if (is_S[i] & (i == 0 || !is_S[i - 1])) {
        is_LMS[i] = true; LMSs.push_back(i);
      }
    }

    // 2. 適当なLMSの順番でinduced sortする(LMSに注目するとLMS-substringの辞書順になっている)
    VI pseudo_sa = induced_sort(str, LMSs, is_S, k);
    // LMSをsubstringの辞書順に並び替える
    VI orderedLMSs(LMSs.size());
    int index = 0;
    for (int x: pseudo_sa) {
      if (is_LMS[x]) { orderedLMSs[index++] = x; }
    }
    // 隣り合うLMS-substringを比較してrankをつけていく
    pseudo_sa[orderedLMSs[0]] = 0;
    int rank = 0;
    if (orderedLMSs.size() > 1) { pseudo_sa[orderedLMSs[1]] = ++rank; }
    REPI (i, 1, orderedLMSs.size() - 1) {
      bool is_diff = false;
      REP (j, n) {
        int p = orderedLMSs[i] + j;
        int q = orderedLMSs[i + 1] + j;
        if (str[p] != str[q] || is_LMS[p] != is_LMS[q]) {
          is_diff = true; break;
        }
        if (j > 0 && is_LMS[p]) { /* assert(is_LMS[q]); */ break; }
      }
      // 同じときは番号をインクリメントしない
      pseudo_sa[orderedLMSs[i + 1]] = is_diff ? ++rank : rank;
    }
    // LMS-substringをひとつのcharとみなしたnew_strを作成
    VI new_str(LMSs.size());
    index = 0;
    REP (i, n) {
      if (is_LMS[i]) { new_str[index++] = pseudo_sa[i]; }
    }

    // 3. 再帰でnew_strのsuffix arrayを求める(これがもとのLMS-substringの辞書順になっている)
    VI LMS_sa;
    if (rank + 1 == LMSs.size()) {
      // LMS-substringの重複がないときは順序そのまま
      LMS_sa = orderedLMSs;
    } else {
      LMS_sa = sa_is(new_str, rank + 1);
      // もとのstrにおけるindexに変換
      for (int& x: LMS_sa) { x = LMSs[x]; }
    }

    // 4. 正しいLMSの順番でinduced sortする(正しいsuffix arrayが得られる)
    return induced_sort(str, LMS_sa, is_S, k);
  }

  // LMS -> L -> M(LMS含む)の順番にbucketに詰めていく。
  // はじめのLMSsがsuffixについてソートされていれば全てのsuffixをソートする。(suffix arrayを返す)
  // はじめのLMSsがソートされていなかった場合、結果はLMSに注目するとLMS-substringでソートされている。
  // k: 現れる文字の種類数
  VI induced_sort(const VI& str, const VI& LMSs, const vector<bool>& is_S, const int k) {
    int n = str.size();
    VI buckets(n); // 結果を格納

    // 出現する文字の個数の累積和(各bucketのはじめのindexを表す)
    VI chars(k + 1);
    for (int c: str) { chars[c + 1]++; }
    REP (i, k) { chars[i + 1] += chars[i]; }

    VI count(k); // 各bucketにいくつ入ったか
    // LMSをbucketに後ろから詰める
    for (int i = LMSs.size() - 1; i >= 0; i--) {
      int c = str[LMSs[i]];
      buckets[chars[c + 1] - 1 - count[c]++] = LMSs[i];
    }

    // 昇順、i - 1 がLであったらbucketに前から詰める
    count = VI(k);
    REP (i, n) {
      if (buckets[i] == 0 || is_S[buckets[i] - 1]) { continue; }
      int c = str[buckets[i] - 1];
      buckets[chars[c] + count[c]++] = buckets[i] - 1;
    }

    // 逆順、i - 1 がSであったらbucketに後ろから詰める
    count = VI(k);
    for (int i = n - 1; i >= 0; i--) {
      if (buckets[i] == 0 || !is_S[buckets[i] - 1]) { continue; }
      int c = str[buckets[i] - 1];
      buckets[chars[c + 1] - 1 - count[c]++] = buckets[i] - 1;
    }
    return buckets;
  }
public:
  string S; int N;
  VI sa; // sa[i]: suffixが辞書順i番目となる開始位置のindex
  SuffixArray(string str_in) : S(str_in), N(str_in.size()) {
    str_in += "$";
    VI str(N + 1);
    REP (i, N + 1) { str[i] = str_in[i] - '$'; }
    sa = sa_is(str, 128);
    sa.erase(sa.begin());
  }
}