ダブリングによる接尾辞配列(Suffix Array)構成 in PyPy3

この記事は[Optimind Advent Calendar 2023](https://qiita.com/advent-calendar/2023/optimind)の23日目の記事となります。


こんにちは。株式会社オプティマインドの最適化チーム所属の伊豆原(イズハラ)と申します。

今回は自社の競プロ部活動の一環として、接尾辞配列というものについて説明したいと思います。

## 接尾辞配列について

文字列に関わるデータ構造で競プロ文脈で頻出するものに**接尾辞配列**(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](https://stackoverflow.com/questions/18495728/suffix-array-using-manber-myers-algorithm)

https://github.com/benfulton/Algorithmic-Alley/blob/master/AlgorithmicAlley/SuffixArrays/sa.py

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

```python
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は既に決定されてます。

```python
    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に突っ込んで次回の判定に回します。

```python
    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)$で計算しますので、まぁ先に言うと今回の実装では勝てません。

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

動かす環境は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を組んだほうが手っ取り早い気もします。

```python
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とどのくらいの速度差になるか気になるところですね。
読んで頂きありがとうございました。