Wavelet MatrixとFM-index

この記事はISer Advent Calendar 2019の1日目として書かれました。

最近ライブラリの剪定をしているときに「そういえばWavelet Matrix実装してないなぁ」と思ったので実装してみたら想像以上にシンプルで驚いたのでそれについての軽い解説を書きました。競プロ方面への応用については色んな方が記事を書いているので、後半ではWavelet Matrixの全文検索への応用としてFM-indexの話をします。

Wavelet Matrixとは
FM-indexとは
感想

Wavelet Matrixとは

データ構造を知る時に最も衝撃を受ける瞬間はどういう操作がどんなオーダーでできるか知る時ですよね。なのでWavelet Matrixができる操作をまずは紹介します。静的な文字列(数列)Sに対して以下のようなクエリに高速に(具体的にはアルファベット集合の大きさの対数時間)答えられるようなデータ構造です。(動的にもできますがここでは静的な場合を扱います。)

クエリ名	操作
access(i)	S[i]
rank(c,i)	S[0,i)中の文字cの数
select(c,i)	Sのi番目(0-indexed)のcの場所
kthMax(l,r,k)	S[l,r)でk番目に大きい文字
rangefreq(l,r,low,high)	S[l,r)のlow以上high未満の文字の数
prevvalue(l,r,low,high)	S[l,r)でhigh未満low以上を満たす文字で最大のもの(すなわちhighのprev value)
topk(l,r,k)	S[l,r)で出現回数が多い文字k個を頻度とともに返す

この表だけでも色んなことができる感じがしますね。この表のクエリ以外にも高速に処理できる操作はたくさんあります。例えばkthMaxを用いれば任意の区間の中央値を高速に求めることができますし、topkを用いれば区間の最頻値も簡単に求まります。次の章からWavelet Matrixの説明をしていますがかなり端折った説明なのでここやここのスライドを見たほうが絶対にわかりやすいと思います。

Wavelet Tree

Wavelet Matrixと同等の時間計算量で同じタイプのクエリを処理できるデータ構造にWavelet Treeがあります。こちらのほうが視覚的にわかりやすいのでまずこちらのイメージをつかんでからWavelet Matrixの説明をすることにします。 Wavelet Treeでは文字列(数列)の値を上位bitから順に0か1かで分類していきます。例として[0,8)の区間の値からなる数列 7232514067125137を例に用います。各bitが0ならば左へ、1ならば右へ値を移動することにすると各段階で数列は下のようになります。

7232514067125137 (3bit目で分類)
2321012137546757 (2bit目で分類)
1011232235457677 (1bit目で分類)
0111222334556777

二分木の構造が見えやすいように間をあけると下のようになります。

7232514067125137 (3bit目で分類)
232101213    7546757 (2bit目で分類)
1011   23223  545  7677 (1bit目で分類)
0 111 222 33 4 55 6 777

Wavelet Treeではこの木の情報を上のような値を直接管理するのではなく各数字が左に行ったか(着目するbitが0だったか)右に行ったかを01で表すことによって保持します。例では以下のようになります。

1000101011001001
1110001011001101
1011010011011011

この01だけから1段下の対応する位置はどうやって求めればよいかを考えてみます。0,1それぞれの場合で、自分より前の0と1が何個あるかがわかれば、行き先のノードのどのインデックスが自分に対応しているかがわかります。例えば、例において添字2(0-indexed)の3は3bit目で分類したあと左の子の添字1に対応する位置に来ますが、これは自分より前に0が1つあることからわかります。

完備辞書(簡潔ビットベクトル)

先程述べたような自分より前の0,1の数を数える、x番目の0,1の添字を求める、といった操作を定数時間で実現するデータ構造が完備辞書です。完備辞書の機能自体は簡単に別の方法で実装できますが、完備辞書は簡潔データ構造と呼ばれるものの一つでその機能を省メモリで実現できることが特徴です。これについての解説はMisterが去年記事を書いてくれているのでこちらに丸投げしようと思います。(投げやり)

misteer.hatenablog.com

重要なことは元データのbit数の線形オーダーのbit数の補助データでビットベクトルのrank,selectが実現できるということです。

Wavelet Treeの計算量

文字列のアルファベットの集合の大きさを $\sigma$ とします。Wavelet Treeではほとんどのクエリは木の深さのオーダーで処理することができ、 $O(log\sigma)$ となります。Wavelet Treeの問題点は空間計算量です。ビット列に $nlog\sigma$ ,補助データに $O(nlog\sigma)$ ,そして木のポインタに $O(\sigma logN)$ 必要になります。しかし $\sigma$ が非常に大きくなるとポインタに必要な空間がネックになってきます。このポインタをなくし、木を辿りながら対応する区間のインデックスを計算するようにしたのがWavelet Matrixです。(Segment Treeにおいて左右の子を2k+1,2k+2とする実装のようなイメージ)

Wavelet Matrix

Wavelet MatrixではWavelet Treeの時と同様に各bitで値を左右に振り分けていきますが、Wavelet Treeでは木構造がそのまま見えます(すなわち2つの部分木の境界を上のような図上で線で引いて表現できる)が、Wavelet Matrixでは各深さのノードの順序が入れ替わっておりその木構造がパッとは見えません。先程と同様に例を通して見てみることにします。各深さにおいて着目するbitが0ならば左1ならば右へ移動するのは変わりませんが、より前の深さでの分類を考えずにWavelet Treeにおける1段目かのようにデータを移動します。

7232514067125137 (3bit目で分類)
2321012137546757 (2bit目で分類)
2322376771011545 (1bit目で分類)
3311155777022264

データの動きとしては基数ソートに似ていますね(基数ソートとの違いは下位bitから見るかどうかです)。実は上のデータ列にはWavelet Treeの場合の各ノードに対応する連続区間があります。それがわかりやすいように空白を入れたものが以下です。

7232514067125137 (3bit目で分類)
232101213 7546757 (2bit目で分類)
23223 7677 1011 545 (1bit目で分類)
33 111 55 777 0 222 6 4

Wavelet Treeの場合の例と見比べると順番は違いますが確かにノードが存在しています(存在は簡単に示せます)。では7546757に対応する頂点から左の子である7677に移動する、といった操作を具体的にどのように計算すればよいかを考えてみることにします。(区間のインデックスから左右の子に対応するインデックスが計算できることでポインタが不要となるのでした)

Wavelet Matrixにおいても以下のようにWavelet Treeと同様に上の例のようなデータを各値が左右どちらに行ったかの01の列として保持します。そして補助データとして各深さにおいて0がいくつあるか(zeroNum)も保持しておきます。

1000101011001001 (zeroNum: 9)
0001110100110010 (zeroNum: 9)
1011001000100010 (zeroNum: 10)

少し考えれば、深さdの[l,r)に対応する頂点の左の子は深さdの列における[rank(0,l),rank(0,r)),右の子は[zeroNum[d]+rank(1,l),zeroNum[d]+rank(1,r))であることがわかります。この式の形からこの計算も完備辞書を用いれば定数時間でできることがわかります。これで頂点のポインタを持たずにWavelet Treeと同等のことができることがわかりました。

各クエリの実装

以上でWavelet Matrixの構築は終わりです。あとは構築したデータを使ってクエリに答えていきます。各クエリはこの木上を辿ったり探索することで木の深さのオーダー $O(log\sigma)$ で処理することができます。実際の細かい手順は先程のリンクのスライドで丁寧に説明されているので今回は詳細な説明は省きますが、実装してみると理解が深まると思います。自分はaccess,rank,select,kth-max(quantile),rangefreq,prevvalue,nextvalueあたりを実装しました。この後紹介するFM-indexの実装もここにあります。

github.com

FM-indexとは

ここからはせっかく作ったWavelet Matrixを使って何かしてみようというお話です。

Suffix Array

文字列検索には様々なアルゴリズムがありますがSuffix Array(接尾辞配列)を構築して二分探索することで長さnの文章Sから長さmのパターンTを $O(mlogn)$ で求めるアルゴリズムがあります。Suffix Arrayというのは接尾辞(その開始インデックスで表す)を辞書順にソートした列で、例えば"abracadabra"という文字列に対しては以下のようになります

i sa[i]
0  11 (空文字列)
1  10 a
2   7 abra
3   0 abracadabra
4   3 acadabra
5   5 adabra
6   8 bra
7   1 bracadabra
8   4 cadabra
9   6 dabra
10  9 ra
11  2 racadabra

Suffix Arrayの構築は愚直に計算して $O(n^{2}logn)$ ,ダブリングを用いる方法で $O(nlog^{2}n)$ ,みんな大好きSA-ISで $O(n)$ でできることは有名ですが、今回はここには深入りしないことにします。接尾辞でソートされているので、検索パターンを接頭辞に持つような接尾辞はSuffix Array上で連続しており、この区間は二分探索によって求めることができます。例えば上の例で"bra"を接頭辞に持つような接尾辞はSuffix Arrayの[6,8)に対応しており、これからSのsa[6]=8文字目と、sa[7]=1文字目からの3文字が"bra"であることがわかります。この二分探索の計算量は1回の文字列比較に $O(m)$ ,比較を $logn$ 回行うので $O(mlogn)$ となります。 FM-indexではこのSuffix Array上での区間を求める操作を $O(mlog\sigma)$ で実現します。 $\sigma$ はアルファベットのサイズで、charの場合は256ですから、これを定数とみなせばパターンの長さの線形時間でできることになります。驚くべきは元の文字列の長さnに依存していないことです。非常に長い文章からパターンを検索する場合に力を発揮しそうですね。Wavelet Matrixを用いるための準備をいくつか行います。

Burrows Wheeler 変換

今後の準備のためにいくつか記号を定義します。sa[i]文字目からの接尾辞の初めの文字をSf[i]と呼ぶことにします(空文字列のときは(任意の文字より若い)ダミー文字とする)。またC[c]をS中に含まれる文字cより小さい文字の数とします。(累積和で簡単に求まります) 文字列SのBurrows Wheeler変換BWTは以下のように定義されます。

bwt[i] = S[sa[i]-1] (sa[i]>0の時)
         ダミー文字 (sa[i]=0の時)

上の例では

i bwt[i] sa[i]
0   a    11 (空文字列)
1   r    10 a
2   d     7 abra
3   $     0 abracadabra
4   r     3 acadabra
5   c     5 adabra
6   a     8 bra
7   a     1 bracadabra
8   a     4 cadabra
9   a     6 dabra
10  b     9 ra
11  b     2 racadabra

となります。 BWTの直感的な意味は、S[i]を後続の接尾辞S[i+1,n)をキーとしてソートしたものです。このBWTは定義より明らかにSfの各文字が1度ずつ現れます。したがってBWTはSfの順列と思うことができますが、BWTにはよい性質があり、BWTのみからSfを復元することができます。この復元をLF mappingと呼びます。

LF mapping

まずBWTの重要な性質として「各文字についてbtwにおける出現順とSfでの出現順は同じになる」というものがあります。例えば上の例でaに番号をつけると、

i bwt[i] sa[i]
0   a5    11 (空文字列)
1   r     10 a5
2   d      7 a4bra5
3   $      0 a1bra2ca3da4bra5
4   r      3 a2ca3da4bra5
5   c      5 a3da4bra5
6   a4     8 bra5
7   a1     1 bra2ca3da4bra5
8   a2     4 ca3da4bra5
9   a3     6 da4bra5
10  b      9 ra5
11  b      2 ra2ca3da4bra5

aはbwt,Sfのどちらでも5,4,1,2,3の順で現れていることがわかります。これが常に成り立つことが辞書順の定義からすこし考えれば示せます。

そして、Sfはソートした接尾辞の先頭文字なので当然ながら辞書順に並んでいます。この2つの性質を用いてBWTからSfを求める方法を考えてみることにします。例えば上の例でa2 = bwt[8] = Sf[?]の?を求めたいとします。Sfは辞書順であったことからaよりも若い文字は全て?より前にあり、この文字数はC[a]=1(ダミー文字のみ)です。逆にaより大きい文字は?より前にありませんから、後は?より前にaがいくつあるかがわかれば良いです。この個数は先程述べた性質よりbwtで8より前にあるaの個数に一致します。やっとWavelet Matrixが使えそうな形が現れてきましたね。bwtからWavelet Matrixを構築すればこの値はrank(a,8)として計算できます。 rank(a,8)=3(a5,a4,a1の3つ)から?=3+1=4と求めることができます。確かにSf[4]=a2となっていることが上の例でも確認できます。ここまでの話を式にまとめると、 c=bwt[i]の時,Sf[rank(c,i)+C[c]]=bwt[i]となります。

FM-index

いよいよLF-mappingを用いて検索パターンTに対応するsaのインデックスの区間[l,r)を求めていきます。空文字列から初めて、T[m-1,m)に対応する範囲、T[m-2,m)に...,T[0,m)に対応する範囲というように更新していきます。空文字列は任意の文字列の接頭辞ですからl=0,r=n+1として初期化します。今T[i+1,m)に対応する区間[l,r)が求まっておりT[i,m)に対応する区間を求めようとしているとしましょう。[l,r)の中でbtw[idx]=T[i]となるようなidxにしか興味がありませんが、そのような最小のidxが存在するならrank(T[i],idx)+C[T[i]]=rank(T[i],l)+C[T[i]]、同様にそのような最大のidxについてもrank(T[i],idx+1)+C[T[i]]=rank(T[i],r)+C[T[i]]なので、

l = rank(T[i],l)+C[T[i]]
r = rank(T[i],r)+C[T[i]]

とするだけでT[i,m)に対応する区間が求まります。

例で確認してみましょう。"abracadabra"から"bra"に対応する[6,8)を求めます。

(再掲)
i bwt[i] sa[i]
0   a    11 (空文字列)
1   r    10 a
2   d     7 abra
3   $     0 abracadabra
4   r     3 acadabra
5   c     5 adabra
6   a     8 bra
7   a     1 bracadabra
8   a     4 cadabra
9   a     6 dabra
10  b     9 ra
11  b     2 racadabra

l = 0, r = 12
[ ""に対応する区間は[0,12) ]
l = rank('a',0)+C['a'] =  0+1 = 1, r = rank('a',12)+C['a'] = 5+1 = 6
[ "a"に対応する区間は[1,6) ]
l = rank('r',1)+C['r'] =  0+10 = 10, r = rank('r',6)+C['r'] = 2+10 = 12
[ "ra"に対応する区間は[10,12) ]
l = rank('b',10)+C['b'] =  0+6 = 6, r = rank('b',12)+C['b'] = 2+6 = 8
[ "bra"に対応する区間は[6,8) ]

ちゃんと求まっていることが確認できますね。

実装

Suffix Array, Wavelet Matrix, Cにあたる配列の計算さえ実装してしまえばたった数行で書けてしまいます。

pair<int,int> searchFMIndex(
    SuffixArray &sa,
    WaveletMatrix<unsigned char,8> &wm,
    vector<int> &lessCount,
    string t){
  int l = 0, r = sa.size()+1; // because of dummy character
  for(int i=0;i<t.size();i++){
    unsigned char c = t[t.size()-1-i];
    l = lessCount[c]+wm.rank(c,l);
    r = lessCount[c]+wm.rank(c,r);
    if(l>=r)return make_pair(-1,-1);
  }
  return make_pair(l,r-1);
}

性能の比較

前にも述べたとおりFM-indexが $O(mlog\sigma)$ 、Suffix Array+二分探索が $O(mlogn)$ なので長い文章から長めの文章を検索すれば差がつくかなと思って実験してみました。初めに100MBの圧縮技術とかのベンチマーク用text fileを落としてきてやろうとしたのですがどうやらメモリとりすぎで怒られてしまうようだったので5MBの聖書を何回かconcatしてギリギリのサイズの文字列を作って、適当な場所から10000000文字くらい取ってきたもので検索してみました。結果が以下です。

Length of the original text 32112346
Suffix array construction finished.
Burrows Wheeler Transformation finished.
Wavelet Matrix construction finished.
Array C construction finished.
Length of the pattern 10000000
[Wavelet Matrix] time: 22302.421999999999[ms]
[Suffix Array] time: 8.371000000000[ms]
[Wavelet Matrix Answer] 17540883,17540887
[Suffix Array Answer] 17540883,17540887
verdict : OK

！！！！！！FM-index遅！！！！！！！

感想

完備辞書の実装が悪いのかなにかわかりませんが色んな長さの文字列を検索してみても全部二分探索のほうが早くて笑いました。所詮log同士の戦いなので劇的に早くはならないのかなとは思っていたのですがここまでとは思いませんでした、定数って厳しいですね。

実はWavelet Matrixの実装は2Aの実験の課題にもなっていましたが簡潔ビットベクトルでちゃんと簡潔性をサボらない実装が面倒で当時はやめてしまいました。もやもやしていたので1年越しにちゃんと実装できてよかったです(結局簡潔性はある程度サボっていますが)。

苦手で避けがちだった文字列が最近個人的に楽しくなってきてPalindromic TreeとかSuffix Automatonの盆栽をそろそろしたいなという気持ちでいます。みなさんもデータ構造の盆栽しましょうね。

明日2日はHelloRuskさんの記事です、お楽しみに。