Suffix Array について

Suffix Array (SA) は文字列探索、全文検索等に用いられるデータ構造です。たまたま AIZU ONLINE JUDGE をもくもく解いている最中にこれを実装したい場面がありました。SA を構築するアルゴリズムとしては SA-IS というものがあり、概要や実装に関してネットで見つかる記事では接尾辞配列 - Shogo Computing Laboratory や Go言語による高速な接尾辞配列の実装 - Qiita がわかりやすいと思います。

ただそれでも SA-IS は個人的に難しく感じ、何故その操作で上手く行くのかピンと来ないという部分も多かったので、実際に論文の証明を読みながら整理してみたいと思いました。

主に参考にした論文は以下の 2 つになります。

また論文の C 実装をほぼそのまま Go に移植しただけですが、SA-IS 実装を sais.go - gist に置いています。

SA を使用した部分一致検索
SA-IS アルゴリズムの概要
バケットで大まかにソートする
2 段階ソートの証明
LMS-suffix のソートに $ S_1 $ の SA が使える理由
LMS-substring のソート方法

1. SA を使用した部分一致検索

例えば $ S = rarb $ という文字列 $ S $ に ar という文字列 $ R $ が含まれているかを調べたいとします。最も単純な方法としては $ S $ の前方から一文字ずつずらしながら $ ar $ と一致するかを見る、というのが考えられます。これは $ S $ の長さを $ n $, $ R $ の長さを $ m $ としたときに $ O((m+n-1)m) $ の時間がかかり、最悪のケースでは $ O(n^2) $ となります。

上の戦法は表現を変えると「 $ S $ の suffix の中に $ R $ を prefix に持つものが存在するか」を見ているともいえます。例の場合、 $ S = rarb $ の suffix には rarb, arb, rb, b が存在し、このうち arb が $ R = ar $ を prefix に持つため、 $ R $ が $ S $ に含まれているとわかります。

ここでもしこの suffix が辞書順にソートされていると、suffix と R との比較回数を最大でも n/2 に減らすことができ、つまり $ O(m \log n) $ の時間で部分一致検索ができます (単純な検索で。実際は Wikipedia によれば $ O(m + \log n) $ の時間で検索できるアルゴリズムがあるそうです)。このとき使用する $ S $ の suffix を辞書順に並べた表を文字列 $ S $ の Suffix Array (SA) と呼びます。

このように SA を使用した検索はわかりやすいし早いしと嬉しいことだらけなのですが、問題はこの SA を作成することにあります。長さ $ n $ の文字列 $ S $ に対する SA を単純なソートアルゴリズムで作成すると $ O(n \times n \log n) = O(n^2 \log n) $ の時間が必要になります。SA を効率良く構築するアルゴリズムは数多く研究されているようなのですが、ここではその中でも SA-IS と呼ばれる手法を見ていきます。この手法は (他のものに比べれば?) 理解しやすく $ O(n) $ の時間で SA を構築することができます。

2. SA-IS アルゴリズムの概要

SA-IS アルゴリズムが登場する論文は Two Efficient Algorithm for Linear Time Suffix Array Construction です。ここではこの論文を (勝手に) SA-IS 論文と呼びます。

SA-IS 論文を読む前にいくつか事前に論文中での表記をまとめておきます。

SA を作成する対象の長さ n の文字列は $ S $ で表されます。 $ S $ は

アルファベットで構成される
n 個の suffix は辞書順で大小が決定される
末尾にどの文字よりも小さい番兵 (‘$’ で表記) がつけられる

という性質を持ちます。

$ S $ の各文字は 0-index の配列表記 $ S[i] $ ( $ 0 \le i \le n-1 $ ) で表します。

$ S[i] $ から始まる suffix は $ suf(S, i) $ と表記します。

これらの表記の簡単な例を挙げると

$ S = rarb$ $ という長さ $ n = 5 $ の文字列について、

$ S[0] = r \\ S[4] = $ \\ suf(S, 1) = arb$ \\ suf(S, 0) \gt suf(S, 4) $

といった感じになります。

論文では Fig. 1 で SA-IS の概要が示されているので、それに沿う形で流れを見ていきます。

各 suffix の S, L-type 分類
LMS-substring のソート
$ S_1 $ の生成
$ SA_1 $ の生成
$ SA_1 $ から $ SA $ の導出

下記の説明では例として $ S = abababab$ $ を使用します。

2-1. 各 suffix の S, L-type 分類

各 suffix, $ suf(S, i) $ は $ suf(S ,i+1) $ との大小関係により S, L-type に分類されます。 $ suf(S, i) < suf(S, i+1) $ のとき、 $ suf(S, i) $ を S-type suffix, $ suf(S, i) > suf(S, i+1) $ のとき、L-type suffix と呼びます。例外として $ suf(S, n-1) $, つまり番兵のみの suffix は S-type と決められます。

例の文字列 $ S $ の場合、

index	0	1	2	3	4	5	6	7	8
S	a	b	a	b	a	b	a	b	$
type	S	L	S	L	S	L	S	L	S

となります。

S, L-type の分類は $ S $ を右 (末尾) から左へ一度走査すれば決定できるので、$ O(n) $ 時間で実行できます。

2-2. LMS-substring のソート

$ SA $ と S, L-type suffix との面白い関係として、「S-type あるいは L-type の全ての suffix をソートできれば、もう片方の type は機械的にソートできる」というのがあります (アルゴリズムに機械的に、という表現が適当かは怪しいんですが)。これがどういうことかというのは 2段階ソートの「ソートをせずにソートする」を見て頂ければと思います。この性質は面白いのですが直感的に納得し辛いとも自分は感じたので、後程この証明を 2 段階ソートの証明で取り上げます。

この手法自体は KA アルゴリズムから来ているのですが、SA-IS では更にソートが必要な対象を狭めて LMS-substring がソートできれば他は機械的にソートができるというように改善されています。

LMS-substring とは、を見る前に LMS-suffix の説明が必要になります。$ suf(S, i) $ が S-type なとき $ suf(S, i-1) $ が L-type であれば $ suf(S, i) $ を LMS-suffix (leftmost S-type suffix) と定義します。

例の文字列 $ S = abababab$ $ の場合、

index	0	1	2	3	4	5	6	7	8
S	a	b	a	b	a	b	a	b	$
type	S	L	S	L	S	L	S	L	S
LMS			*		*		*		*

で $ suf(S, 2), suf(S, 4), suf(S, 6), suf(S, 8) $ が LMS-suffix です。

LMS-substring は $ S[i..j] $ なる文字列のうち、

(1) $ i \ne j $ かつ $ suf(S, i), suf(S, j) $ の両者が LMS-suffix かつ $ suf(S, k) $ ( ただし $ i < k < j $ ) なる LMS-suffix が存在しない

(2) $ S[n-1] $, つまり番兵のみの LMS-suffix

のいずれかを満たす文字列のことと定義します。

例の文字列 $ S = abababab$ $ の場合、

$ S[2..4] = aba \\ S[4..6] = aba \\ S[6..8] = ab$ \\ S[8] = $ $

が LMS-substring になります。

SA-IS でソートが必要になる LMS-substring の数は元の文字列 $ S $ の長さ $ n $ の半分以下であり (SA-IS 論文の Lemma 3.5)、しかも 2段階ソートと同様な方法でソートができるのですが、その手順と何故それが上手くいくのかについては後程 LMS-substring のソート方法で見ていきます。

(ちなみに 2 段階ソートというのは 2 段階ソートに倣って勝手にそうよんでいるのですが、論文中の用語だと induced sort というのがそれに当たるのだと思います)

2-3. $ S_1 $ の生成

2-2 では 2 段階ソートという手法と LMS-substring をソートすればいいということを紹介しました。しかしなぜ LMS-substring をソートすることで 2 段階ソートが利用できるのかというところはまだ曖昧なままです。

2 段階ソートは全ての S-type suffix がソートできれば L-type suffix のソートを可能にしますが、この前提はもう少し緩めて LMS-suffix だけにすることができます。というのは 2 段階ソートの処理対象は「 $ suf(S, i) $ が S-type で $ suf(S, i-1) $ が L-type」なので LMS ではない S-type suffix は条件を満たさないためです。

そして LMS-suffix の大小関係なのですが、これを SA-IS では $ S $ から生成した $ S_1 $ という新たな文字列の SA ( $ SA_1 $ と定義する) によって求めることができます。

$ S_1 $ は少し説明が難しいのですが自分は「LMS-substring のソート後に小さい順に 0 から始まる番号をつけ、それを $ S $ における出現順に並べて作成する文字列」と理解しています。

$ SA_1 $ が何故 LMS-suffix のソートに使用できるかの解説としては SA-IS の「LMS-substring」がわかりやすいと思います。ただ自分はこれでもうまく理解できなかったので、後程 LMS-suffix のソートに $ S_1 $ の SA が使える理由で証明も見てみます。

例の $ S $ では

LMS-substring	辞書順	S における出現順
S[2..4] = aba	2	0
S[4..6] = aba	2	1
S[6..8] = ab$	1	2
S[8] = $	0	3

なので $ S_1 = [2, 2, 1, 0] $ となります。substring が同じ文字列になる場合、同じ数字が割り当てられる (ここでいう aba に対する 2) という点に注意です。

2-4. $ SA_1 $ の生成

前述のように $ SA_1 $ は $ S_1 $ の SA です。論文中の SA-IS 実装例では 2-3 で生成された $ S_1 $ に応じて直接 $ S_1 $ から $ SA $ を生成する場合と SA-IS のアルゴリズムを $ S_1 $ に再帰的に適用する場合とに処理が分岐するのですが、ここでは詳細を省略します。

2-5. $ SA_1 $ から $ SA $ の導出

$ SA_1 $ から目的の $ SA $ を導出する処理の詳細は SA-IS 論文中では 3.3 で取り上げられています。手順の概要は

$ SA_1 $ を利用してソートした LMS-suffix を後方から順に、 $ SA $ の適当なバケットの後方から詰めていく
$ SA $ を前からスキャン。$ SA[i] $ が存在し、かつ $ suf(S, SA[i]-1) $ が L-type ならば $ SA[i] - 1 $ を適当なバケットの前方から詰めていく
- L-type suffix のソートが完了
$ SA $ を後ろからスキャン。$ SA[i] $ が存在し、かつ $ suf(S, SA[i]-1) $ が S-type ならば $ SA[i] -1 $ を適当なバケットの後方から詰めていく
- S-type suffix のソートが完了

となります。この手順でソートできることは後程見る 2 段階ソートの証明でわかることなのでここでは詳細を省略します。

またここでバケットという用語を出したのですが、これは SA を suffix の頭一文字でグループ化したもので、バケットで大まかにソートするでもう少し見ていきます。

以上で SA-IS アルゴリズムの概要をおさえたのですが、所々直感的に理解し辛い箇所があったのでその部分について論文中の証明を見て考えてみます。

3. バケットで大まかにソートする

文字列 $ S $ の全 suffix を早く正確にソートするのは難しいのですが、各 suffix がソート後に大体どのあたりに来るかはバケットソートでわかります。SA-IS では出現しうる文字 (ここではアルファベット) 毎にバケットを用意し、各 suffix の先頭一文字を見て適当なバケットに入れることで suffix が SA 上どの index の範囲に登場し得るかという情報を利用します。

例えば $ S = abababab$ $ の場合、suffix のうち ‘$’ で始まるものが 1, ‘a’ で始まるものが 4, ‘b’ で始まるものが 4 なので SA と含まれるバケットは以下のようになります。この表を見ると ‘a’ で始まる suffix であれば必ず $ suf(S, i) $ (ただし $ 1 \le i \le 4 $) であることがわかります。

No.	bucket	suffix
0	$
1	a
2	a
3	a
4	a
5	b
6	b
7	b
8	b

SA-IS におけるバケットソートの有り難い点として、各バケット内では L-type suffix が必ず S-type suffix よりも前に来る、というのがあります。これはバケット内では先頭の文字が一致していること、 S-type, L-type の定義より導けます。証明は KA アルゴリズムの Lemma 2.2 です。

簡単な例を挙げると、cbcd$ という文字列中の cbcd$ (L-type), cd$ (S-type) という suffix を比較した場合、先頭一文字が一致しているので 2 文字目以降の bcd$, d$ という文字列で比較することになります。このとき L-type であれば 2 文字目は 1 文字目と同じかそれよりも大きい文字 (ここでは b) であり、S-type であれば同じか小さい文字 (ここでは d) になります。なので両 suffix が先頭一文字 (ここでは c) と同じでない場合は、直感的に S, L-type というだけで大小の比較ができることがわかります。先頭一文字と同じ文字の場合も後続の文字を見ていけば必ずどこかで同様の比較が行われることになります。

4. 2 段階ソートの証明

2 段階ソートというのは正式な名称では無いと思うのですが、KA アルゴリズムの Lemma 2.3 で証明の対象になっている操作のことです。

操作の概要は

$ S $ の全 suffix をバケットソートする
予めソートした S-type suffix を大きい順に該当するバケットの最後尾に詰めていく
- S-type のソートが完了
$ SA $ を前からスキャンし、$ SA[i] $ が存在し、かつ $ suf(S, SA[i]-1) $ が L-type の場合に、該当するバケットに前から詰めていく
- L-type のソートが完了

です。2 段階ソートでは図付きで解説されているのでそちらの方がわかりやすいと思います。

ステップ 2 についてはバケットで大まかにソートするで紹介した「バケット内では L-type の suffix が S-type よりも前に来る」という SA と S, L-type の性質から導けます。

論文ではステップ 3 の操作が正しいことを $ SA[i] $ のスキャン時にそこに存在する $ suf(S, SA[i]) $ がソートされた正しい位置にあるという性質を数学的帰納法で証明することで確認しています。

証明の流れとしてはステップ 3 で $ SA[i] $ に達したとき、もし $ SA[i] $ が正しい位置にいない (つまり辞書順で i 番目の suffix ではない) ならば、 $ i < k $ なる $ k $ が存在し $ SA[k] $ がそこにいるべきなのですが、そんな $ k $ は存在しないということを背理法で示しているようです。

ポイントは

(a) $ SA $ はバケットでまとめられているので、 $ k $ が $ i $ と違うバケットから来ることはない
- 順番が前後するならそれはバケット内になる
(b) $ suf(S, k) $, $ suf(S, i) $ は両方とも L-type
- S-type はステップ 1, 2 より既にソート済みなのでこのような状況に陥ることはない

です。

(a) より $ suf(S, i) $, $ suf(S, k) $ の先頭一文字 ( $ c $ とする) は一致するので

$ suf(S, i) = c \alpha \\ suf(S, k) = c \beta $

と表すことができ、また順番が前後しているという仮定より $ \beta < \alpha $ となります。

一方 (b) より $ SA[k] $ は L-type なので

$ \beta < c \beta = SA[k] \tag 1 $

$ \beta $ も $ S $ の suffix の 1 つであり、(1) と数学的帰納法の仮定より $ \beta $ は $ SA[0] から SA[i-1] $ のどこかに出現しているはずです。そうなると $ \beta < \alpha $ なのでステップ 3 の操作的に $ SA[k] $ は $ SA[i] $ よりも前に出現するはずで、あとに出現するはずが無いことがわかります。

5. LMS-suffix のソートに $ S_1 $ の SA が使える理由

証明は SA-IS 論文の Lemma 3.8 になります。

まず証明中に出てくる $ P_1, S_1 $ ですが、これらはそれぞれ「LMS-suffix の $ S $ 中での index を並べた配列」, 「LMS-substring 同士を比較して順序をつけ、それを $ S $ における出現順に並べて作成する文字列」になります。個人的に $ S_1 $ がとてもわかりにくいと感じたのですが、例を挙げると $ S = abababab$ $ について

index	0	1	2	3	4	5	6	7	8
S	a	b	a	b	a	b	a	b	$
type	S	L	S	L	S	L	S	L	S
LMS			*		*		*		*

なので $ P_1 = [2, 4, 6, 8] $, $ S_1 = [2, 2, 1, 0] $ となります。

証明で述べられているのは、$ S_1 $ の suffix 間の比較がそのまま LMS-suffix 間の比較に使えるということになります。例えば $ suf(S_1, 0) > suf(S_1, 3) $ から $ suf(S, P_1[0]) > suf(S, P_1[3]) $ が導けます。

証明は $ S_1 $ がユニークな数字で構成されるか否かで場合分けされます。ここでユニークというのは、n 個の LMS がある場合に $ S_1 $ に 0 から n-1 の数字が一回ずつ出現するということです。

ユニークな数字で構成される場合は、各 LMS-substring がそこから始まる LMS-suffix の prefix であり、LMS-substring の辞書順で LMS-suffix 間の順序も決定できるということを考えると明らかです。

そうでない場合として例えば上の $ suf(S, 2), suf(S, 4) $ の比較を考えてみます。辞書順による比較という観点では $ suf(S, 2) $ は index 2, 4, 6, 8 の LMS-substring を繋げたもの、$ suf(S, 4) $ は同様に 4, 6, 8 の LMS-substring を繋げたものと見ることができます。このとき繋ぎ目間の文字は重複しますが、これは順序の比較には影響しません。それぞれ対応する $ S_1 $ の suffix は 2210, 210 であり、この比較が元の文字列 $ S $ の suffix の比較に代用できます。

6. LMS-substring のソート方法

SA-IS で最終的にソートしなければならない対象は LMS-substring であり、これをいかに効率良くソートできるかが問題になります。SA-IS 論文の 3.4 で $ O(n) $ でソートを行うアルゴリズムが示されています。

LMS-suffix を $ SA $ バケットの後ろから詰めていく
$ SA $ を前からスキャンし、 $ SA[i] $ が存在し、かつ $ suf(S, SA[i]-1) $ が L-type ならば $ SA[i] - 1 $ を適当なバケットの前方から詰めていく
$ SA $ を後ろからスキャンし、$ SA[i] $ が存在し、かつ $ suf(S, SA[i]-1) $ が S-type ならば $ SA[i] -1 $ を適当なバケットの後方から詰めていく

論文ではこの手順で LMS-substring がソートできることを示すために $ pre(S, i) $ と表記される LMS-prefix を定義しています。 $ pre(S, i) $ は (1) $ i = n-1 $ (番兵) ならば番兵自身、(2) そうでなければ i 以降最初に出てくる LMS の index を j として $ S[i..j] $ の文字列となります。 $ S = abab$ $ を例に取ると以下のようになります。

index	0	1	2	3	4
S	a	b	a	b	$
type	S	L	S	L	S
LMS			*		*

$ pre(S, 0) = S[0..2] = aba \\ pre(S, 1) = S[1..2] = ba \\ pre(S, 2) = S[2..4] = ab$ \\ pre(S, 3) = S[3..4] = b$ \\ pre(S, 4) = S[4] = $ $

また $ suf(S, i) $ の S, L-type に応じて $ pre(S, i) $ の type も定義されるものとします。

Theorem 3.12 では上のソート操作により LMS-substring ではなく LMS-prefix がソートされることを示します。といっても上の表を見るとわかるように LMS-prefix は全ての LMS-subsring を含むので、LMS-prefix がソートできれば結果的に LMS-substring もソートできることになります。

まずステップ 1 では LMS-suffix を $ SA $ に詰めていきますが、基本的に正しい順序で入る保証はありません。ただ例外が一つあり、番兵のみの suffix である $ suf(S, n-1) = $ $ についてはバケット ‘$’ のただ一つの要素であることが保証されるため、常に正しい位置 $ SA[0] $ に入ります。

次にステップ 2 では L-type LMS-prefix がソートされることを数学的帰納法で示します。はじめて入れる L-type LMS-prefix については、既に $ SA $ に含まれている suffix が S-type のみなので入れ方的に必ずソートされた位置に入ります。

$ k > 1 $ 個の L-type LMS-prefix がソートされた位置に入ったとして、 $ k + 1 $ 個目の L-type LMS-suffix である $ pre(S, i) $ について考えます。背理法で示すために、仮に別の L-type LMS-prefix である $ pre(S, j) > pre(S, i) $ が既に同一バケットに入っているとします。このとき

$ S[i] = S[j] $ (同一バケットより)

$ pre(S, j+1) > pre(S, i+1) $ ( 同一バケットかつ $pre(S, j) > pre(S, i) $ より)

となります。

ここで $ pre(S, j) $ が $ pre(S, i) $ より前に入っているので、ステップ 2 で行う操作を踏まえると、$ SA $ 内で $ pre (S, j+1) $ は $ pre(S, i+1) $ でより前に出現しなければならないことがわかります。しかしそうすると数学的帰納法の仮定よりこの 2 つはソートされていなければならないのに $ pre(S, j+1) > pre(S, i+1) $ であることと矛盾します。従ってそのような $ pre(S, j) $ は存在せず、ステップ 2 により L-type LMS-prefix が正しくソートできることが示されます。

(と論文だとこのように証明されているように読めるのですが、ここで $ pre(S, j+1), pre(S, i+1) $ がソートされていなければならない、と言えるのかで悩みました。この両者が L-type suffix であれば数学的帰納法の仮定よりそう言えますし、S, L-type suffix の組み合わせの場合もバケットと S, L-type の性質よりそうだろうと言えそうです。しかし LMS-suffix 同士の場合はステップ 1 の手順で入れられているだけなのでソートされている保証がなさそうに思います。恐らく LMS-suffix 同士でもバケットが異なれば問題無く、またバケットが同じ場合は LMS-prefix で見れば $ pre(S, i) = pre(S, j) $ となり前提を満たさなくなるので ok… という理解でいいんでしょうか)

ステップ 3 により全 S-type LMS-prefix がソートされますが、これはこれまで見た内容と重なる部分も多いので省略します。

以上によりステップ 1, 2, 3 で LMS-prefix がソートされ、結果として LMS-substring がソートできるということがわかりました。

Suffix Array について

1. SA を使用した部分一致検索

2. SA-IS アルゴリズムの概要

2-1. 各 suffix の S, L-type 分類

2-2. LMS-substring のソート

2-3. \( S_1 \) の生成

2-4. \( SA_1 \) の生成

2-5. \( SA_1 \) から \( SA \) の導出

3. バケットで大まかにソートする

4. 2 段階ソートの証明

5. LMS-suffix のソートに \( S_1 \) の SA が使える理由

6. LMS-substring のソート方法