Suffix Array と LCP と文字列検索の実装をした

この土日のメモです. SAとLCPのお気持ちをまとめたくなっただけ. 間違ってたらごめん

文字列アルゴの勉強する気が起きないたった一つの理由: Rolling Hash
— νιυεζ (@xiuez) 2019年12月13日

これをやめたいので, 手始めにSuffix Arrayを使った文字列検索をやってみようかなというのが今回の主題

概要

SA-ISでSuffix Arrayを構築 $O(|S|)$
LCP配列の構築 $O(|S|)$
LCPによるSuffix同士のLCPをSparse Tableで構築 $O(|S| \log{|S|})$ , クエリ $O(1)$
Suffix Arrayの二分探索で文字列検索を $O(|T| log{|S|})$
Suffix ArrayとLCPの二分探索で文字列検索を構築 $O(|S|)$ , クエリ $O(|T| + \log{|S|})$

の実装をやってみました. このときのメモを残しておきたいと思います.

各用語の説明はここではしません... 他の記事や, 蟻本 - Amazonを参考に.

SA-IS

Suffix Arrayの実装は蟻本にも載っていますが, そこまで早くありません... SA-ISというアルゴリズムが早いらしいのでこれを実装します.

SA-ISの理解には, この記事がとても参考になりました. とてもわかりやすい記事です.

mametter.hatenablog.com

SA-ISの実装はyosupoさんのコードを見ました.

僕が書いたSA-ISのコードはこれです.

judge.yosupo.jp

メモリを使い回す(resizeの回数を減らしてメモリを使いまわしても, assignが割と早くてこれが非自明)
push_backをなくす
入出力の早いライブラリを使うともっと早くなります

実はSA-ISの論文に実装が載っていてそれがとても速いです. ぜひ参考にしてみてください.

以下, 文字列SのSuffix ArrayをSAとします. Sの辞書順でi番目に小さいsuffixをSuf[i] := S[SA[i]...]とします.

LCP(Longest Common Prefix)

LCP配列は, Suffix Arrayで隣り合ったSuffix(つまり, Suf[i]とSuf[i + 1])の最長共通接頭辞を求めた配列です. Kasai's Algorithmを用いて $O(|S|)$ で構築できます.

LCPの理解は以下の記事がわかりやすいです. 蟻本にもあるはず.

lumakernel.github.io

例は, 接尾辞配列(Suffix-Array) | Luzhiled’s memo がわかりやすいです.

僕の実装は先頭に無(空配列)があるので, 以下のようになります.

i :lcp
0 : 0 
1 : 0 a
2 : 1 abra
3 : 4 abracadabra
4 : 1 acadabra
5 : 1 adabra
6 : 0 bra
7 : 3 bracadabra
8 : 0 cadabra
9 : 0 dabra
10: 0 ra
11: 2 racadabra

任意のsuffix同士のLCP

上の例で, i = 2, abraと, j = 5, adabraのLCPを求めるとすると, 3, 4, 5のlcpの最小値である1がその答えになります.

Suffix Arrayで, indexが $i$ のsuffixと $j$ のsuffixのLCPは, $[i + 1, j + 1)$ 間のlcpの最小値になります.

なので, lcpをSparse Tableに載せると構築 $O(|S| \log{|S|})$ , クエリ $O(1)$ で処理できます.

Suffix Arrayで文字列検索

文字列SのSuffix Array SAを使って, Sの中に文字列Tがあるかどうかを二分探索で処理できます. これは, Suffix Arrayによって各suffixがソートされているのを利用しています.

計算量は $O(|T| \log{|S|})$ です. AOJの提出コード

cin >> t;
int L = 0;
int R = sa.size();
while(R - L > 1) {
  int M = (L + R) >> 1;
  if(s.substr(sa[M], t.size()) <= t) {
    L = M;
  }
  else {
    R = M;
  }
}
cout << (s.substr(sa[L], t.size()) == t) << endl;

これがかなりはやいなんでだろう

SAとLCPで文字列検索

この二分探索はさらに高速化できます. suffixとTの比較を最小限にすることで, $O(|T| + \log{|S|})$ を達成します.

具体的には, suf[L]とTのLCPを常に持ちながら二分探索をします. このLCPをLlcpとします. M = (L + R) / 2として, suf[L]とsuf[M]のLCPを求めて, nlcpとします. nlcpは先に書いたとおり, Sparse Tableで求めることができます. 次にLlcpとnlcpを比較します.

Llcp < nlcpのとき

以下の例で考えてみます. (Suffix Arrayではありませんが, 複数の文字列を辞書順にソートしたという意味で同じです)

T = ad

L : aaa
    aaab
    aaac
M : aac
    aacc
    ba
R : 

Llcp = LCP(aaa, ad) = 1  // "a"aa, "a"dなので
nlcp = LCP(aaa, aac) = 2 // "aa"a, "aa"cなので

Tは辞書順でsuf[T]以上ということがわかっているので, Llcp < nlcpより, Tとsuf[M]のLCPはLlcpであり, Tは辞書順でsuf[M]以上です. なので, Llcpはそのままで, L = Mとします.

Llcp > nlcpのとき

T = aaac

L : aaa
    aaab
    aaac
M : aac
    aacc
    ba
R : 

Llcp = LCP(aaa, aaac) = 3  // "aaa", "aaa"cなので
nlcp = LCP(aaa, aac) = 2 // "aa"a, "aa"cなので

Tとsuf[M]のLCPはnlcpであり,Tは辞書順でsuf[M]未満です. なので, Llcpはそのままで,R = M`とします.

Llcp = nlcpのとき

T = aacc

L : aaa
    aaab
    aaac
M : aac
    aacc
    ba
R : 

Llcp = LCP(aaa, aacc) = 2  // "aa"a, "aa"ccなので
nlcp = LCP(aaa, aac) = 2 // "aa"a, "aa"cなので

このときは, Tとsuf[M]の辞書順の関係がわからないので比較をします. このとき, LCPの部分は一致していることがわかっているので比較をしなくてよいです. 比較をした後, Llcpを比較をした時の計算結果を利用して更新します.

Llcpは探索中, 単調増加します. なので, 文字列の比較が全体で $O(|T|)$ しかされません. これにより, 計算量が改善されます.

実際にコードを示します.

std::pair<int, int> get_lcp(const std::vector<T>& t, int si, int offset) {
  int i = offset;
  si += offset;
  while(i < t.size() && si < N) {
    if(t[i] != str[si]) {
      return { i, t[i] - str[si] };
    }
    i++;
    si++;
  }
  return { i, 0 };
}

std::pair<int, int> search(const std::vector<T>& t) {
  int L = 0;
  int R = N + 1;
  int Llcp = 0;

  while(R - L > 1) {
    int M = (L + R) >> 1;
    int nlcp = st.query(L + 1, M + 1);
    if(Llcp < nlcp) {
      L = M;
    }
    else if(Llcp > nlcp) {
      R = M;
    }
    else {
      auto p = get_lcp(t, sa[M], Llcp);
      if(p.second >= 0) {
        L = M;
        Llcp = p.first;
      }
      else if(p.second < 0) {
        R = M;
      }
    }
  }

  return { Llcp, L };
}

これで早くなるはず...!

onlinejudge.u-aizu.ac.jp

5倍遅くなった...

Sparse Tableの構築が重すぎる

$O(|S| \log{|S|})$ 流石に重い... 改善したい

Sparse Tableを使わない方法で改善

二分探索だけならSparse Tableである必要はありません. Segment Treeを使います.
二分探索で最小値を求めたい区間は必ず[L, (L + R) / 2)に対応できます. なので, 二分探索するときに, Segment Treeのノードを降りていくようにすると構築 $O(|S|)$ で二分探索ができるようになります.

コードはこんな感じ

...
...
  seg_n = 1;
  while(seg_n < N + 1) seg_n <<= 1;
  seg.resize(seg_n * 2, 1e9);
  for(int i = 0;i + 1 < N + 1;i++) {
    seg[i + seg_n - 1] = lcp[i + 1];
  }
  for(int i = seg_n - 1; i --> 0;) {
    seg[i] = std::min(seg[(i << 1) + 1], seg[(i << 1) + 2]);
  }
}

std::pair<int, int> get_lcp(const std::vector<T>& t, int sa_i, int offset) {
  if(sa_i > N) return { offset, -1 };
  int i = offset;
  int si = sa[sa_i] + offset;
  while(i < t.size() && si < N) {
    if(t[i] != str[si]) {
      return { i, t[i] - str[si] };
    }
    i++;
    si++;
  }
  return { i, 1 };
}

std::pair<int, int> search(const std::vector<T>& t) {
  int L = 0;
  int R = seg_n;
  int Llcp = 0;
  int j = 0;

  while(R - L > 1) {
    int M = (L + R) >> 1;
    int nlcp = seg[(j << 1) + 1];
    if(nlcp == 1e9) {
      j = (j << 1) + 1;
      R = M;
    }
    else if(Llcp < nlcp) {
      j = (j << 1) + 2;
      L = M;
    }
    else if(Llcp > nlcp) {
      j = (j << 1) + 1;
      R = M;
    }
    else {
      auto p = get_lcp(t, M, Llcp);
      if(p.second >= 0) {
        j = (j << 1) + 2;
        L = M;
        Llcp = p.first;
      }
      else if(p.second < 0) {
        j = (j << 1) + 1;
        R = M;
      }
    }
  }

  return { Llcp, L };
}

onlinejudge.u-aizu.ac.jp

これでも最初の二分探索に勝てませんでした... なんでだろうでもこれでもかなり速いです.

しめ

FM-indexとかやってみたくなりました