ダブリングによる接尾辞配列（Suffix-Array）の実装と個人的なメモ

AC Libraryに、より高速なライブラリが存在する。

https://atcoder.github.io/ac-library/document_ja/string.html

以下に載せた実装は、自分の理解を深めるために実装したもの。（蟻本のものにアレンジを加えclass化した。計算量は変わらず、Manber & Myers の O(n log² n) 時間のアルゴリズム）

前提

接尾辞配列（Suffix-Array）とは、ある文字列の接尾辞（Suffix）を辞書順にソートしたもの。

ナイーブにSuffix-Arrayを構築しようとすると、文字列をsとして、空間計算量 |s|^2、時間計算 |s|^2 log |s| となる。後述するが、Suffix-Arrayを構築するのは文字列中に含まれる部分文字列の検索を高速に行うためであり、ナイーブなSuffix-Arrayの構築の時間計算量は、ナイーブな検索の時間計算量(|s|^2)を上回ってしまうので、このままでは構築する意味があまりない。また空間計算量も、扱う文字列の長さが10⁵以上の問題ではMLEになってしまう。

そのため、Suffix-Array構築時の計算量を改善する必要がある。

Suffix-Arrayでもつ値は部分文字列ではなくSuffixの始点のindexにする。(これで空間計算量は|s|になる)

各indexからのSuffixの辞書順の順位が求められれば、その順位を使ってSuffix-Arrayのindexをソートすることができるので、rank配列を作り、これをダブリングによって段階的に更新していく。

以下は、rank配列を更新する実装部分についての簡単な説明。

rank配列の初期化

rankはそのindexの文字コードによって初期化している。rank[n](空文字)だけは、-1で初期化。

    for (int i = 0; i <= n; i++) {
      suffix_array[i] = i;
      rank[i] = i < n ? (int)s[i] : -1;
    }

確定した2*k文字まで長さのrankを元にSuffix-Arrayをソートする

k=1から始めて、各indexから2*k文字まで見た時の順位をrank配列に格納していく。

nを対象の文字列の長さとして、2*k>=nまで更新を行えば全てのSuffixのrankを求めたことになる。

ループ回数が log n 回、各シーケンス内のソートが n log n 時間なので、時間計算量は n log² n

    vector<int> tmp(n + 1); // 一度tmp配列に
    for (k = 1; k <= n; k *= 2) {
      sort(suffix_array.begin(), suffix_array.end(), compare);
      tmp[suffix_array[0]] = 0;
      for (int i = 1; i <= n; i++) {
        tmp[suffix_array[i]] = tmp[suffix_array[i - 1]] + (compare(suffix_array[i - 1], suffix_array[i]) ? 1 : 0);
      }
      for (int i = 0; i <= n; i++) {
        rank[i] = tmp[i];
      }
    }

ソートに使う比較関数

k=1から始めて、2*kまでの長さの部分文字列の大小比較を行う。

    auto compare = [&](int a, int b) {
      if (rank[a] != rank[b]) return rank[a] < rank[b]; // 前半k文字のrankが異なるなら、後半k文字を比較するまでもなく大小が確定する
      else {
        // 前半k文字のrankが同等なら、後半のk文字のrankの大小で2*k文字の大小が決まる
        // +k文字目が、文字列をはみ出すなら空文字としてrank=-1とする
        auto rank_ak = a + k <= n ? rank[a + k] : -1;
        auto rank_bk = b + k <= n ? rank[b + k] : -1;
        return rank_ak < rank_bk;
      }
    };

実装

class SuffixArray {
private:
  string s;
  int n, k;
  vector<int> suffix_array;
  vector<int> rank;
  void create_suffix_array() {
    for (int i = 0; i <= n; i++) {
      suffix_array[i] = i;
      rank[i] = i < n ? (int)s[i] : -1;
    }
    auto compare = [&](int a, int b) {
      if (rank[a] != rank[b]) return rank[a] < rank[b];
      else {
        auto rank_ak = a + k <= n ? rank[a + k] : -1;
        auto rank_bk = b + k <= n ? rank[b + k] : -1;
        return rank_ak < rank_bk;
      }
    };
    vector<int> tmp(n + 1);
    for (k = 1; k <= n; k *= 2) {
      sort(suffix_array.begin(), suffix_array.end(), compare);
      tmp[suffix_array[0]] = 0;
      for (int i = 1; i <= n; i++) {
        tmp[suffix_array[i]] = tmp[suffix_array[i - 1]] + (compare(suffix_array[i - 1], suffix_array[i]) ? 1 : 0);
      }
      for (int i = 0; i <= n; i++) {
        rank[i] = tmp[i];
      }
    }
  }

public:
  SuffixArray(const string &_s) {
    s = _s;
    n = _s.size();
    suffix_array.resize(n + 1);
    rank.resize(n + 1);
    create_suffix_array();
  }

  // 文字列sの部分文字列として、パラメータで渡した文字列tが現れるかを返す
  bool contains(const string &t) {
    int l = 0, r = n;
    while (r - l > 1) {
      int mid = (l + r) / 2;
      int index = suffix_array[mid];
      if (s.compare(index, t.size(), t) < 0) l = mid;
      else r = mid;
    }
    return s.compare(suffix_array[r], t.size(), t) == 0;
  }
};

Suffix-Arrayを使った文字列検索

上記の実装のcontainsという関数は、文字列sの中にtが現れるかを判定している。

Suffix-Arrayがあることで二分探索によって部分文字列の検索を行えるので、O(|t| log |s|) 時間で判定を行える。

検証用問題

Suffix-Arrayを使った高速文字列検索によって解くことができる。（上記のcontainsにより検証可能）

onlinejudge.u-aizu.ac.jp

Baby Stepsなブログ

競プロとか。間違ったこと書いてあったら@pi0nep1oneにご連絡ください。