接尾辞配列について

文字列に関わるデータ構造で競プロ文脈で頻出するものに接尾辞配列(Suffix Array)というものがあります。長さ$N$の文字列$S$の部分文字列の集合$\{S[i:N]|0\leq i \leq N\}$を辞書順でソート時の開始位置$i$のなす配列です。

例: "abracatabra"

i	S[i:N]	rank
0	abracatabra	3
1	bracatabra	7
2	racatabra	10
3	acatabra	4
4	catabra	8
5	atabra	5
6	tabra	11
7	abra	2
8	bra	6
9	ra	9
10	a	1
11	(empty)	0

=> suffix array : [11, 10, 7, 0, 3, 5, 8, 1, 4, 9, 2, 6]
(空文字を入れるかは流儀に依る気がします)

二分探索による高速な文字列検索や、最長共通接頭辞(Longest Common Prefix)の検索に応用があるようです。

構成法

接尾辞配列の構成方法についてSA-IS法というものがあります。これは$O(N)$で構成できる非常に高速なものですが、文字ごとにタイプを分けて管理するなどなかなか難しいアルゴリズムかと思います。中身を理解し実装するのは個人的に骨が折れます。

一方、Manber & Myersに(多分)基づくダブリングによる構成法があります。計算量は$O(N\log(N)^2)$ですが理解し易いと個人的に感じますので、本記事ではこちらの紹介をさせていただきます。

実装元となっているのは(記憶が合ってれば……)StackOverFlowの以下の質問で紹介されている以下のGitHubにあるソースコードだったかと思います。ただしこれは、質問のコメントにあるようにダブリングが行われていない実装のようでして、これをダブリングするように修正します。
StackOverFlow - suffix array using manber myers algorithm

実装の全体は以下になります。以下、こちらを説明します。

from collections import defaultdict, deque
def ManberMyers(S, conv=ord):
    N = len(S)
    d = defaultdict(list)
    for i,s in enumerate(S):
        d[conv(s)].append(i)
    rank = [0]*(N+1)
    i = 1 # offset for empty
    tbd = deque([])
    for k in sorted(d.keys()):
        for t in d[k]:
            rank[t] = i
        i += len(d[k])
        if len(d[k]) > 1:
            tbd.append((d[k],1))
    while tbd:
        target, l = tbd.popleft()
        d = defaultdict(list)
        for t in target:
            nt = rank[t+l] if t+l < N else -1
            d[nt].append(t)
        i = 0
        for k in sorted(d.keys()):
            for t in d[k]:
                rank[t] += i
            i += len(d[k])
            if len(d[k]) > 1:
                tbd.append((d[k],l*2))
    result = [0]*(N+1)
    for i,r in enumerate(rank):
        result[r] = i
    return result

挙動を"abracatabra"で説明します。

まず初期準備として最初のrankを設定します。変数rankには$rank[i] = S[i:N]$が"取りうる最高順位"が入ります。計算の中でこれを更新していき、最終的に正しい順位に収めていきます。例えば'a'から始まる部分文字列は必ず1〜5位のどれかになるので、最初は1を設定します。'b'なら6〜8位なので6です。
なお、この時点でrankが一意に定まる(=文字列中で唯一の文字)ならばそのインデックスのrankは既に決定されてます。

    for i,s in enumerate(S):
        d[conv(s)].append(i)
    rank = [0]*(N+1)
    i = 1 # offset for empty
    tbd = deque([])
    for k in sorted(d.keys()):
        for t in d[k]:
            rank[t] = i # 可能な最高順位を設定
        i += len(d[k])
        if len(d[k]) > 1: # 決定したインデックスは処理しない
            tbd.append((d[k],1))

a  b  r  a  c  a   t  a  b  r  a  *
1  6  9  1  8  1  11  1  6  9  1  0

辞書dにはキーを順位付けとし、値に対応するインデックスのリストが入ります。順位付けには初回は文字のorder、後のwhileループではそこまでで决定されているrankを利用します。
キューである変数tbdには同じ文字の成すインデックスのリストと1のペアが入ります。キューに追加するリストは今後も、その時点で同じrankを有するインデックスのリストが入り、数字にはその時点で判明している一致する文字列長が入ります。

d = {97: [0, 3, 5, 7, 10], 98: [1, 8], 114: [2, 9], 99: [4], 116: [6]}
tdb = deque([([0, 3, 5, 7, 10], 1), ([1, 8], 1), ([2, 9], 1)])

次にwhileループに入ります。

変数targetはその時点でのrankが同じインデックスのリストで、長さlの分だけ対応する部分文字列が一致することが約束されています。l+1文字目以降の情報でrankをより細かく順序付けるのですが、そこでそれまでに計算されたrankの値を利用します(ダブリング)。ただし文字列長を超えるならば、辞書順で最も小さいと捉えて-1を設定します。

rank[t+l]が一致し順序付け出来なかったインデックスについては長さl*2分一致してることになりますので、tbdに突っ込んで次回の判定に回します。

    while tbd:
        target, l = tbd.popleft()
        d = defaultdict(list)
        for t in target:
            nt = rank[t+l] if t+l < N else -1 # ダブリングによる順序付け
            d[nt].append(t)
        i = 0
        for k in sorted(d.keys()):
            for t in d[k]:
                rank[t] += i
            i += len(d[k])
            if len(d[k]) > 1: # ランクが決定できなかったら次に回す
                tbd.append((d[k],l*2))

target = [0, 3, 5, 7, 10], l = 1 # (tbdから先頭をpop)
d = {6: [0, 7], 8: [3], 11: [5], -1: [10]}
rank = [2, 6, 9, 4, 8, 5, 11, 2, 6, 9, 1, 0] # 1,4,5位は決定。2位はまだ候補が2つある([0,7])。
tbd = deque([([1, 8], 1), ([2, 9], 1), ([0, 7], 2)]) # 未決定の[0,7]をキューに追加

あとはtbdが空になるまでループを回せば、すべての順位付けが完了した状態になります。

……ちなみに挙動から分かる通り、変数tbdに入るlは1->2->4->...と遷移していくのでリストを2個用意してfor-loopの形にすることも可能です。多分そのほうがちょっと速いかも？まぁ好みで。

速度計測

実装できたので早速、速度を測りたいと思います。比較対象はAtCoder社様からコンテスト中に提供されているライブラリAtCoder Libraryからstring.suffix_arrayを使用します。これは計算量$\mathcal{O}(N)$で計算しますので、まぁ先に言うと今回の実装では勝てません。

動かす環境はMacBookProでプロセッサは2GHz クアッドコアIntel Core i5。メモリは16GBです。

なお、あまりに遅いのも悔しいので、後で付記するちょっと速くした版(Manber Myers(Fast))も測ってます。先程のリスト化や実装中のtargetの長さが2の時のアンロールなどを行っています。
計測用の文字列は以下を用意しました(接尾辞配列のパフォーマンスチェックにはどういう文字列が適切なんですかね？)

random: ランダムなアルファベットの文字列
all 'a': すべての文字が'a'の文字列
repeat 2: 長さ$N/2$のランダムなアルファベットの文字列を2回繰り返したもの

case: $N=2 \times 10^5$

	random	all 'a'	repeat 2
AtCoder Library	110ms	19ms	73ms
Manber Myers	119ms	180ms	905ms
Manber Myers(Fast)	105ms	94ms	184ms

case: $N=10^6$

	random	all 'a'	repeat 2
AtCoder Library	540ms	98ms	854ms
Manber Myers	551ms	1338ms	5900ms
Manber Myers(Fast)	538ms	1396ms	1362ms

こう見ますと、randomのときはいい勝負してますがall 'a'のときは圧倒的に負けてますね。実際、all 'a'では実装の中では長さ$10^6$に近い配列を何回も作る形になるのでだいぶ遅くなります。できる限り作られた配列を使い回すような処理を入れればマシになると思いますが……

おまけ

計測に使用したManber Myers(Fast)の実装です。工夫すれば色々速くなると思いますが、それならSA_ISを組んだほうが手っ取り早い気もします。

def ManberMyersFast(S, conv=ord):
    N = len(S)
    d = defaultdict(list)
    for i,s in enumerate(S):
        d[conv(s)].append(i)
    rank = [0]*(N+1)
    i = 1 # offset for empty
    tbd = [] # use list
    for k in sorted(d.keys()):
        for t in d[k]:
            rank[t] = i
        i += len(d[k])
        if len(d[k]) > 1:
            tbd.append(d[k])
    ntbd = []
    l = 1
    while tbd:
        for target in tbd:
            if len(target) == 2: # unroll
                x,y = target
                nx = rank[x+l] if x+l < N else -1
                ny = rank[y+l] if y+l < N else -1
                if nx==ny:
                    ntbd.append(target)
                elif nx < ny:
                    rank[y] += 1
                else:
                    rank[x] += 1
                continue
            d = defaultdict(list)
            for t in target:
                nt = rank[t+l] if t+l < N else -1
                d[nt].append(t)
            i = 0
            for k in sorted(d.keys()):
                for t in d[k]:
                    rank[t] += i
                i += len(d[k])
                if len(d[k]) > 1:
                    ntbd.append(d[k])
        tbd, ntbd = ntbd, []
        l *= 2
    result = [0]*(N+1)
    for i,r in enumerate(rank):
        result[r] = i
    return result

まとめ

以上、ダブリングによる接尾辞配列作成でした。C++とかで組むとSA_ISとどのくらいの速度差になるか気になるところですね。
読んで頂きありがとうございました。