足跡-sokuseki-

りかの日進月歩の記録

Suffix ArrayのO(NlogN)構築（とおまけでSA-IS）

競技プログラミング

この記事はCompetitive Programming (1) Advent Calendar 2018 - Adventar の25日目の記事です。

はじめに

競技プログラミングの問題でSuffix Arrayを使う場合、蟻本で紹介されている O(N log^2 N)の構築で間に合う問題がほとんどです。しかし、 O(N log^2 N)では定数倍高速化をしないとTLEしてしまう問題もあるらしい*1ので、より高速な構築アルゴリズムを紹介したいと思います。

Suffix Arrayの定義

Suffix Arrayを日本語にすると「接尾辞配列」です。
長さ n の文字列 s の接尾辞は全部で n 個あります。
たとえば、 s = abracadabra の接尾辞は以下の11個あります。文字列の左にある数は、それぞれの接尾辞のインデックス（文字列 s の何文字目以降か）を表しています。

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

この n 個の接尾辞をソートした時に、それぞれの接尾辞が辞書順で何番目であるかを配列にしたものが接尾辞配列になります。
文字列 s の接尾辞配列を SA とすると
SA[ i ] = 辞書順で i 番目に小さい接尾辞のインデックス
となります。

たとえば s = abracadabra の接尾辞配列は以下のようになります。

i SA[i] 接尾辞
0 10 a
1 7 abra
2 0 abracadabra
3 3 acadabra
4 5 adabra
5 8 bra
6 1 bracadabra
7 4 cadabra
8 6 dabra
9 9 ra
10 2 racadabra

接尾辞配列があると文字列検索などが高速でできます。

アルゴリズム

接尾辞のソートと巡回シフト

厳密に言うと、紹介するアルゴリズムは接尾辞をソートするのではなく、文字列の巡回シフトをソートするのですが、これは文字列の末尾に記号$を追加すれば実現できます。ここで、$はどのアルファベットよりも辞書順で小さいことにします。
ソートされた巡回シフトの順序はソートされた接尾辞の順序と等しくなります（接尾辞は巡回シフトにおける$以降を無視したものと同じなので）。

例えば、 dabbb の巡回シフトと接尾辞をソートしたものは以下のようになります。

1.  abbb$d  abbb
4.  b$dabb  b
3.  bb$dab  bb
2.  bbb$da  bbb
0.  dabbb$  dabbb

巡回シフトをソートするために、巡回部分文字列について考えます。
s の部分文字列についてs[ i … j ]という表記をします。ここで i > j のときは文字列s[ i … n-1 ]と文字列s[ 0 … j ]を連結したものを表します。
（以下では適宜、文字列のインデックスはmod n の値と考えてください）

このアルゴリズムは \lceil log n \rceil + 1 回の反復が行われます。 k 番目（ k = 0 … \lceil log n \rceil ）の反復では、 s の長さ 2^kの巡回部分文字列 n 個をソートします。最後の \lceil log n \rceil 番目の反復では、長さ2^{\lceil log n \rceil}の部分文字列がソートされます。 n \le 2^{\lceil log n \rceil}より、これは巡回シフト全体をソートすることと同じです。

巡回シフトのソートに使用する配列

p[ i ] を、長さが 2^k である部分文字列のうち、辞書順で i 番目となる部分文字列のインデックスとします。
例えば s = aaba は k = 2 のとき p = (3, 0, 1, 2) となります。

アルゴリズムの各反復では、順列p[ ]に加えて、配列c[ ]も保持します。c[ i ]は i 番目の部分文字列が属する同値類に相当します。この同値類は0から始まる数で管理され、ある文字列が別の文字列より辞書順で小さい場合は同値類のラベルも小さくなり、2つの文字列が等しい場合は等しくなります。
なぜ同値類を意識する必要があるかというと、部分文字列の中には全く同じ文字列も含まれていて、アルゴリズムはそれらを等しく扱う必要があるからです。
以下のソースコードでは、同値類の数はclassesという変数に保存されます。

s = aabaのとき、それぞれの反復における部分文字列、p[ ]、c[ ]は以下のようになります。

k = 0 : (a,a,b,a)
        p = (0,1,3,2), c = (0,0,1,0)

k = 1 : (aa,ab,ba,aa)
        p = (0,3,1,2), c = (0,1,2,0)

k = 2 : (aaba,abaa,baaa,aaab)
        p = (3,0,1,2), c = (1,2,3,0)

重要な点は、p[ ]の値がすべて異なるということです。
また、たとえば、 k = 0 のとき、p[ ]は(3, 1, 0, 2)でも(3, 0, 1, 2)でも良いです。等しい部分文字列のp[ ]の値は入れ替わっても問題ありません。等しい部分文字列であるかどうかはc[ ]が管理しているからです。

巡回シフトのソートアルゴリズム

文字列 s を引数にして、巡回シフトをソートした順列を返す関数sort_cyclic_shiftsを考えます。
最初（ k=0 のとき）は長さ1の部分文字列をソートします。長さ1の部分文字列の種類は多くないので、カウントソートを使うと O(N) で実装できます。カウントソートでは、各文字について文字列に何回出現するかを数え、この情報を使ってp[ ]を構築します。その後、p[ ]を使ってc[ ]を構築します。

vector<int> sort_cyclic_shifts(string const& s) {
    int n = s.size();
    const int alphabet = 256;
    vector<int> p(n), c(n), cnt(max(alphabet, n), 0);

    //k=0のソート

    for (int i = 0; i < n; i++){
        cnt[s[i]]++;
    }
    for (int i = 1; i < alphabet; i++){
        cnt[i] += cnt[i-1];
    }
    for (int i = 0; i < n; i++){
        p[--cnt[s[i]]] = i;
    }
    c[p[0]] = 0;
    int classes = 1;
    for (int i = 1; i < n; i++) {
        if (s[p[i]] != s[p[i-1]]){
            classes++;
        }
        c[p[i]] = classes - 1;
    }

    //以下k=1…logn+1の反復をする（省略）

}

k-1 回目の反復をしてp[ ]とc[ ]が求められたとき、 O(N) で k 回目の配列の値を求めたいです。これが実現できれば、反復は \lceil log n \rceil + 1 回実行されるので全体の計算量は O(NlogN) になります。

これを行う際の重要な点は、長さ 2^k の巡回部分文字列は、長さ 2^{k-1} の2つの巡回部分文字列から構成されるということです。そして、前のステップで得た長さ 2^{k-1} の巡回部分文字列に対応する配列c[ ]の値の情報を使うと、長さ 2^k の巡回部分文字列をO(1)で比較することができます。
したがって、位置 i と位置 j で始まる長さ 2^k の2つの部分文字列を比較するときは (c[i], c[i+2^{k−1}]) のペアと (c[j], c[j+2^{k−1}]) のペアを比較すれば良いです。

f:id:wk1080id:20181221153344p:plain

では実際に k 回目の配列の値はどうやって求めるといいのでしょうか。これは非常に簡単な解決法があり、長さ 2^k の部分文字列をこれらの数値のペアでソートします。これにより、次のp[ ]を得ることができます。ただし、通常のソートを使うと O(NlogN) になってしまうので工夫が必要です。ペアの要素は高々 n なので、再びカウントソートをすることができますが、カウントソートを用いるよりも効率的な方法が存在します。

ここでは基数ソートの基礎となるテクニックを使います。
ペアのソートをするために、まずペアの2つ目の要素に着目してペアを並べ替え、その後で1つ目の要素を使って並べ替えます。（1つ目の要素のソートは安定ソートでなければいけません。安定ソートだと等しい要素の相対的な順序を変えることなくソートできるので、初めに行った2つ目の要素でのソートの順番を生かすことができます。）
しかし、ペアの2つ目の要素は k-1 回目のステップですでにソートされていました。したがって、2つ目の要素によってペアをソートするためには、p[ ]の値から 2^{k-1}を引くだけで良いです。（長さ 2^{k-1} の辞書順最小の部分文字列が位置 i からはじまるとき、長さ 2^k の部分文字列の後半の辞書順最小の部分文字列が始まる位置は i - 2^{k-1} になります。）
よって、単純な減算だけでp[ ]におけるペアの2つ目の要素をソートできます。次に1つ目の要素によって安定ソートを実行する必要があります。これはカウントソートで実現できます。


例えば s = aaba の k = 2 のときを考えます。 k = 1 のとき p[ ]とc[ ]はそれぞれ(0,3,1,2), (0,1,2,0)だったとします。位置 i から始まる長さ4の巡回部分文字列は (c[i], c[i+2]) のペアを見ればよいので、

i ペア 対応する部分文字列
0 (0,2) (aa, ba)
1 (1,0) (ab,aa)
2 (2,0) (ba,aa)
3 (0,1) (aa,ab)

となります。
2つ目の要素によってペアをソートするためには、p[ ]の値から 2^{k-1} = 2 を引けば良いので p = (0,3,1,2) は (2,1,3,0) になります。よって i = 2, 1, 3, 0 について (c[i], c[i+2]) のペアを見ると

i ペア 対応する部分文字列
2 (2,0) (ba,aa)
1 (1,0) (ab,aa)
3 (0,1) (aa,ab)
0 (0,2) (aa, ba)

となり、2つ目の要素でソートできました。あとは1つ目の要素を使ってカウントソートをすれば、ペアのソートが完了します。

p[ ]が計算できたので、残るはc[ ]ですが、これは k=0 のときと同様に、p[ ]を順に調べて隣接するペアを比較すれば計算できます。

以下は残りの実装です。一時的な配列pn[ ]とcn[ ]を使用して、2つ目の要素と新しい同値類c[ ]のインデックスの順列を保存しておきます。

vector<int> sort_cyclic_shifts(string const& s) {
    int n = s.size();
    const int alphabet = 256;
    vector<int> p(n), c(n), cnt(max(alphabet, n), 0);

    //k=0のソート

    for (int i = 0; i < n; i++){
        cnt[s[i]]++;
    }
    for (int i = 1; i < alphabet; i++){
        cnt[i] += cnt[i-1];
    }
    for (int i = 0; i < n; i++){
        p[--cnt[s[i]]] = i;
    }
    c[p[0]] = 0;
    int classes = 1;
    for (int i = 1; i < n; i++) {
        if (s[p[i]] != s[p[i-1]]){
            classes++;
        }
        c[p[i]] = classes - 1;
    }


    //以下k=1…logn+1の反復をする
    vector<int> pn(n), cn(n);
    for (int k = 0; (1 << k) < n; ++k) {
        for (int i = 0; i < n; i++) {
            pn[i] = p[i] - (1 << k);
            if (pn[i] < 0){
                pn[i] += n;
            }
        }
        fill(cnt.begin(), cnt.begin() + classes, 0);
        for (int i = 0; i < n; i++){
            cnt[c[pn[i]]]++;
        }
        for (int i = 1; i < classes; i++){
            cnt[i] += cnt[i-1];
        }
        for (int i = n-1; i >= 0; i--){
            p[--cnt[c[pn[i]]]] = pn[i];
        }
        cn[p[0]] = 0;
        classes = 1;
        for (int i = 1; i < n; i++) {
            pair<int, int> cur = {c[p[i]], c[(p[i] + (1 << k)) % n]};
            pair<int, int> prev = {c[p[i-1]], c[(p[i-1] + (1 << k)) % n]};
            if (cur != prev){
                ++classes;
            }
            cn[p[i]] = classes - 1;
        }
        c.swap(cn);
    }
    return p;
}
最後の処理

前述の通り、このアルゴリズムは接尾辞をソートするのではなく、文字列の巡回シフトをソートします。そのため、与えられた文字列の末尾に、どのアルファベットよりも辞書順で小さい記号$を追加しています。なので、巡回シフトをソートし終わった順列の先頭は、$から始まる巡回シフトになっています。つまり、ソートが終わった後の順列の先頭を最後に取り除く必要があります。それが終われば接尾辞配列の構築は終了になります。

vector<int> suffix_array_construction(string s) {
    s += "$";
    vector<int> sorted_shifts = sort_cyclic_shifts(s);
    sorted_shifts.erase(sorted_shifts.begin());
    return sorted_shifts;
}

計算量について

このアルゴリズムは時間計算量 O(NlogN) 、空間計算量 O(N) です。文字列に使われる文字が少なければ少ないほど（たとえばアルファベットの小文字しか使われないなど）、計算量は小さくなります。

実装例

実装量は蟻本のやつとあんまり変わらないです。

参考文献

この記事の内容は以下のページの日本語訳+αです。構築後に問題に応用する部分の説明も詳しく載っています。また、codeforcesなどの問題も載っています。
Suffix Array - Competitive Programming Algorithms
このページには接尾辞配列をつくるときの文字列には$がないのでこの記事でもそういう説明にしていますが、蟻本とかその他の書籍を見ると大体の場合$も接尾辞配列に含まれています。どっちがいいのかよくわかってないです。

最後に

今日12/25はアドベントカレンダー最終日かつクリスマスです。
それなのにほぼ英語を日本語に直しただけの記事で終わりにするのは少し心苦しいので、昔につくったSA-IS*2の解説スライドを手直ししたやつも貼っておきます。ぜひ読んでください。
2019/3/15追記：スライドの一部が間違っていたため修正しました。
sais.pdf - Google ドライブ
それではみなさん、素敵なクリスマスと素敵な年末をお過ごしください。

*1:A - IOIOIが険しいらしい

*2:接尾辞配列を O(N) で構築するアルゴリズム