概要
wavelet matrix を勉強したので、まとめます。本記事の目標は、merge sort tree（マージソート木）とwavelet tree/matrix（ウェーブレット行列/ウェーブレット木）を領域木の枠組みから統一的に眺めることです。特に wavelet matrix を次のように理解します。
- wavelet matrix は値を分割するタイプの領域木と思える
- 言い換えると、binary trie の各ノードで添字列を管理するのと同じ
- ただし、子の順番をうまく入れ替え、各段をビットベクトル 1 つで管理する
- ビットベクトルの rank 関数が、fractional cascading による索引の役割を果たす
この理解はもちろん既知のもので、そもそも wavelet tree/matrix の論文自体や、書籍『簡潔データ構造』、有名なスライド『ウェーブレット木の世界』などで本質的には同様の説明がなされています。ですが、競プロ er 向けに明示的にそう書いてある記事が意外と少なく、ここでまとめることにしました。
本記事では、wavelet tree/matrix において重要な「簡潔データ構造」の側面をあえて脇に置いているので、その部分は他の記事で補ってください。
Merge sort tree
Merge Sort Tree （マージソート木）は、マージソートの過程で得られる部分的な（各区間に対応する）ソート列を持つデータ構造です。典型的にはセグメント木の各ノードに「その区間に含まれる値の昇順列」を格納します。たとえば数列 $(0,5,2,1,7,1,2,6)$ に対しては次のデータを持ちます。仕切り線に挟まれている列がソートされていることに注意してください。
Merge sort tree 対応する添字の区間
|0 1 1 2 2 5 6 7| [0,8)
|0 1 2 5|1 2 6 7| [0,4) [4,8)
|0 5|1 2|1 7|2 6| [0,2) [2,4) [4,6) [6,8)
|0|5|2|1|7|1|2|6| [0,1)[1,2)[2,3)[3,4)[4,5)[5,6)[6,7)[7,8)
この merge sort tree を使うと、数列 $A = (A_1, \dots, A_N)$ に対して、
- $L,R,X$ が与えられる。$A_L, \dots, A_{R-1}$ のうち値が $X$ 以下となる要素の個数を答えよ
というクエリに $O(\log^2 N)$ で答えられます。具体的には、
- 添字の区間 $[L,R)$ を、セグメント木で管理する $O(\log N)$ 個の区間に分解する
- それぞれの区間に対応するノードが管理するソート列に対して二分探索を行う
とすればよいです。 $O(\log N)$ 個の列にそれぞれ二分探索を行えば、全体の計算量は $O(\log^2 N)$ です。
この「たくさんのソート列上で同じ値を二分探索する」という状況を高速化するのが、次に説明する fractional cascading と呼ばれるテクニックです。
各区間に対応するノードでソート列の累積和も一緒に管理すれば、「$L,R,X$ が与えられるので、$A_L, \dots, A_{R-1}$ のうち値が $X$ 以下となる要素の総和を答えよ」というクエリにも答えられます。
Fractional cascading
一般的な fractional cascading では、次の問題を考えます。
- 最大長 $N$ の、$k$ 個の（ソート済み）配列 $A_1,\dots,A_k$ が与えられる
- 同じ値 $x$ について、$k$ 個の配列それぞれの「$x$ 以下の要素の個数」を答えよ
もちろん各配列ごとに二分探索を行えば、計算量 $O(k \log N)$ でこの問題が解けます。
Fractional cascading は、配列間に補助的な索引を張る前計算を行うことで、
- 最初の 1 回だけ二分探索を行い
- 残りは対応する位置を定数時間でたどる
とすることで $O(\log N + k)$ 時間で問題に答えるテクニックです。
ソート列のマージにおける fractional cascading
今回の目的は merge sort tree への応用なので、fractional cacading の特殊ケースとして、まず次のシンプルな問題を説明します。
- ソート済み配列 $A,B$ が与えられる
- ソート済み配列 $C$ を、 $A$ と $B$ をマージしたソート列とする (つまり
C = sorted(A+B)) - 値 $x$ について、 $C$ で二分探索するだけで、$A,B$ 上での二分探索の結果が手に入るか？
この問題はそこまで難しくはないです。$C$ 上のランク関数 rankA, rankB を
- $\mathrm{rankA}(v)$: $C_0,\dots,C_{v-1}$ のうち $A$ 由来のものの個数
- $\mathrm{rankB}(v)$: $C_0,\dots,C_{v-1}$ のうち $B$ 由来のものの個数
と定めます。すると、 $C$ 上で値 $x$ を二分探索したときの位置を $v$ としたとき（つまり bisect_left(C,x)) 次が成り立ちます。
bisect_left(A,x) = rankA(v)bisect_left(B,x) = rankB(v)
つまり関数 rankA, rankB を前計算しておくことで、$C$ 上で二分探索すれば、$A,B$ では二分探索することなくその位置が求まります。つまりランク関数が索引の役割を果たします。なお$\mathrm{rankA}(v) + \mathrm{rankB}(v) = v$が成り立つので、どちらか一つを前計算しておけば十分です。
たとえば $A=(0,1,2,5)$, $B=(1,2,6,7)$ のとき、$C=(0,1,1,2,2,5,6,7)$ のランク関数は次のようになります（値 $1,2$ は$A,B$ 両方にありますが、$C$ に現れる最初の要素が $A$ 由来とします）。
| C | 0 | 1 | 1 | 2 | 2 | 5 | 6 | 7 | --- |
|---|---|---|---|---|---|---|---|---|---|
| 由来 | A | A | B | A | B | A | B | B | --- |
| $v$ | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 |
rankA(v) |
0 | 1 | 2 | 2 | 3 | 3 | 4 | 4 | 4 |
rankB(v) |
0 | 0 | 0 | 1 | 1 | 2 | 2 | 3 | 4 |
（値 $x$ によらず）二分探索の位置 $v$ だけで $C$ から $A$ への索引が張れるのは、$C$ に $A$ の要素がすべて含まれているからです。一般的な問題設定（最大長 $N$ の数列 $k$ 個）では要素の包含はないので、直接この手法を適用できません。すべての列を結合させた列を考えて、すべての列間に完全な索引を張ることはできますが、前計算や空間計算量が重すぎます。そこで、本家 fractional cascading は次のアイデアで賢く索引を張ることで、前計算を $O(kN)$ に抑えつつ、クエリの計算量 $O(\log N + k)$ ですべての二分探索の位置を計算します。詳細は英wikipediaを参照してください。
- 前の数列 $A_{i+1}$ の要素の偶数番目だけ数列 $A_{i}$に加えてソートする、を $i=k-1,\dots,1$ で繰り返す
- すると $A_i$ の中に $A_{i+1}$ の要素の偶数番目があるので、そこに索引を張れば誤差 $1$ で対応する位置が求まる
- さらに数列の長さの総合計も高々$2$倍に収まる
Merge sort tree との関係
この構造を merge sort tree に応用しましょう。先ほどの説明の通り、$x$ 以下の要素数を数えるクエリでは、
- 同じ値 $x$ に対して
- 複数のソート済み配列で二分探索
します。merge sort tree の各ノードは $2$ つのソート列をマージしたソート列なので、「ソート列の fractional cascading」のアイデアで二分探索の回数を節約できます。具体的には、すべてのノードについてランク関数を前計算しておきます（ソート列をマージするとき同時に計算できるので、この前計算で計算量オーダーは悪化しません）。そうすれば、セグメント木の頂上のノード（すべての値のソート列）で値 $x$ に関する二分探索をすると、ランク関数をたどることで下位ノード上での二分探索の位置が分かります。つまり二分探索の回数が $O(\log N)$ 回からたったの $1$ 回になりました！よってクエリの計算量が $O(\log^2 N)$ から $O(\log N)$ に改善しました。
この視点は、後で述べる wavelet matrix において、ビットベクトルの rank 関数が果たしている役割を理解する上で重要になります。
領域木
（広義の）領域木とは、平面上（一般には $d$ 次元空間上）のクエリに答えるためのデータ構造で、平衡二分木やセグメント木の各ノードで、別の構造（配列・平衡二分木・セグメント木など）を管理するものの総称です。
私の理解では、領域木は結局次のようなものです。
- $x$ 座標を管理する木を作る。木のノードは、$x$ 軸上の区間に対応する
- 各ノードには、$x$ 座標が管理する区間に含まれるような点の情報を適切な形で持つ
- 狭義の領域木 (range tree) では、点を平衡二分木で管理する
計算幾何で言う狭義の領域木 (range tree) では、$d$ 次元空間内の $N$ 点が与えられたとき、
- 「直方体領域 $\prod_{i}[L_i,R_i]$ にある点をすべて探索せよ」
というクエリに $O(\log^{d} N + k)$ （または fractional cascading を使えば $O(\log^{d-1} N + k)$）で答えるものです（$k$ は領域内にある点の個数）。
Merge sort tree が領域木の一例と思えることを確認しましょう。復習ですが、Merge sort tree は、次のように情報を管理します。
- $x$ 座標を管理する木を作る。木のノードは、$x$ 軸上の区間に対応する
- 各ノードは、$x$ 座標が管理する区間に含まれる点の $y$ 座標のソート列を持っている
さて、先ほどのクエリ
- $L,R,X$ が与えられる。$A_L, \dots, A_{R-1}$ のうち値が $X$ 以下となる要素の個数を答えよ
を考えましょう。ここで数列 $A$ の各要素 $A_i$ を、$(i,A_i)$ という平面上の点に対応させます。するとクエリは
- $L,R,X$ が与えられる。長方形領域 $[L,R) \times (-\infty, X]$ に含まれる点の個数を求めよ
という問題に言い換えられます。つまり、先ほど説明した内容は「長方形クエリを merge sort tree で解ける」と言い換えることもできます。
Wavelet tree, あるいは binary trie に添字のソート列を乗せる
Merge sort tree では、添字の区間を分割し、各ノードで値のソート列を管理していました。長方形クエリとして考えるなら、その逆、すなわち
- 値の範囲を分割し
- 各ノードで添字のソート列を管理
するのも自然です。ここで数列の値は $0$ 以上 $2^B$ 未満の整数、つまり $B$ ビット整数であると仮定しましょう。このとき、値を $2$ 分割することは、値を $2$ 進表現で見たときの上位ビットから順に分岐することに対応します。よって、深さ $B$ の binary trie の各ノードに「その値域に属する要素の添字列（昇順）」を保持するデータ構造と考えられます。
この構造が wavelet tree と呼ばれるものです（再注：この説明は簡潔データ構造の側面を無視しています）。この見方をすると、wavelet tree は領域木の一種であり、軸を入れ替えた merge sort tree と考えることもできます。
例えば、数列 $A = (0,5,2,1,7,1,2,6)$ に対して操作を行うと、次のようになります。
昇順の添字列 対応する値域
[0 1 2 3 4 5 6 7] [0,8)
[0 2 3 5 6] [1 4 7] [0,4) [4,8)
[0 3 5] [2 6] [1] [4 7] [0,2) [2,4) [4,6) [6,8)
[0] [3 5] [2 6] [] [] [1] [7] [4] [0,1)[1,2) [2,3)[3,4) [4,5)[5,6) [6,7)[7,8)
Binary trie より直接的に merge sort tree と結びつけるなら、（値、添字）のペアをソートして得られる列の添字の部分だけを抜き出して、その添字列の merge sort tree を作る、と考えることもできます。
順序統計量クエリ
binary trie に昇順の添字列を乗せるタイプの領域木を使うと、
-
range_freq(L,R,x): 区間 $[L,R)$ に含まれる要素の中で$X$ 以下の個数
という先ほどのクエリを同様に解くことができます。つまり値域 $[0,X]$ を binary trie のノードで管理される 値域 $O(\log B)$ 個に分割し、それぞれのノードが管理するソートされた添字列について $[L,R)$ に含まれる添字の個数を数えればよいです。
さらに順序統計量クエリ、すなわち
-
kth_smallest(L,R,k): 区間 $[L,R)$ に含まれる要素の中で$k$ 番目に小さい値
に対応できます。Fenwick 木で二分探索を行うのと同じ感覚です。つまり、例えば値域 $[0,8)$ に対応するノード上にいるとして、次のようにどちらの子ノードに移るかを決めていきます。
- 値域 $[0,4)$ 上で条件をみたす添字が $k$ 個以上なら、求める値は $[0,4)$ にいるので、そのノードに移る
- そうでないなら、求める値は $[4,8)$ にいるので、そのノードに移る
この上の２つのクエリを組み合わせれば、次のクエリにも答えられます。
-
prev_value(L,R,x): 区間 $[L,R)$ に含まれる要素の中で$x$ 未満のもののうち、最大のもの
以上のクエリは、fractional cascading と組み合わせれば、クエリ当たりの計算量は、 $O(\log B)$ です（$B$ は値のビット長でした）。ちなみに、添字列の場合は最初の二分探索すら必要ありません（最初の添字列は$(0,\dots,N-1)$なので、位置が明らかなので）。
Wavelet matrix
概要
Wavelet matrix は、wavelet tree を配列ベースで巧妙に整理した実装と捉えることができます。対応できるクエリは wavelet tree と同じで、次のようなクエリです。
-
range_freq(L,R,x): 区間 $[L,R)$ に含まれる要素の中で$X$ 以下の個数 -
kth_smallest(L,R,k): 区間 $[L,R)$ に含まれる要素の中で$k$ 番目に小さい値 -
prev_value(L,R,x): 区間 $[L,R)$ に含まれる要素の中で$x$ 未満のもののうち、最大のもの
内部構造の理解としては、次の二つがポイントです。
- binary trie の各深さごとに処理をまとめる
- 子の順番をうまく入れ替え、各段をビットベクトル 1 本で管理する
まずはビットベクトルについて簡単に説明して、その後で具体的な wavelet matrix の構造を説明します。
ビットベクトル
ビットベクトルは、0/1 からなる長さ $N$ のビット列に対して、$O(1)$ で次の操作を行なえるデータ構造です。
-
access(i)：$i$ 番目のビットを取得する -
rank0(i)：区間 $[0,i)$ に含まれる値 $0$ の個数を求める -
rank1(i)：区間 $[0,i)$ に含まれる値 $1$ の個数を求める
もちろん累積和でこのクエリに答えることはできますが、そのためには（例えば）$32$ bit 整数を $N$ 個持つ必要があります。そうではなく、$N + o(N)$ bit だけの情報でこのクエリに答えられるデータ構造が（簡潔）ビットベクトル、もしくは完備辞書と呼ばれるものです。割愛した select(i) という操作も含め、簡潔性の意味や具体的な実装は他の記事を参照してください。
Wavelet matrix の構築
$B$ ビット整数の列 $A=(A_1,\dots,A_N)$ が与えられたとき、wavelet matrix を次のように構築します。
- 数列の値を上位ビットから見ていき、次の操作を行う
- そのビットが $0$ か $1$ かだけを記録したビットベクトルを各段に作る
- ビットが $0$ の要素が前に、$1$ の要素が後ろに来るように、数列を安定に並べ替える（同じビットの要素では順番を入れ替えないようにソートする）
- 最終的に、得られた各段のビットベクトルだけをデータとして保持する
具体例で確認しましょう。$A = (0,5,2,1,7,1,2,6)$ に対して、ビット数 $3$ の wavelet matrix を構築します。
- 先頭ビット
- $A$ の各要素の先頭ビットからなるビットベクトルは $01001001$
- この値について$A$をソートして、$A = (0,2,1,1,2,5,7,6)$ と置き換える
- 中央ビット
- $A$ の各要素の中央ビットからなるビットベクトルは $01001011$
- この値について$A$をソートして、$A = (0,1,1,5,2,2,7,6)$ と置き換える
- 末尾ビット
- $A$ の各要素の末尾ビットからなるビットベクトルは $01001001$
- （この値について$A$をソートして、$A = (0,2,2,6,1,1,5,7)$ と置き換える）
- 3 つのビットベクトルを wavelet matrix のデータとして保持する
この構成を wavelet tree と結びつけるため、数列を区切って値域と対応させた図を示します。縦線で区切られている場所は、同じ値域であることを表しています。ただし実際に縦線の位置を覚えるわけではなく、単に列を管理します。
数列と値域 ビットベクトル
値域| 0~7 |
数列|0 5 2 1 7 1 2 6| 01001001 （左の数列の第 2 bit）
添字|0 1 2 3 4 5 6 7|
値域| 0~3 | 4~7 |
数列|0 2 1 1 2|5 7 6| 01001011 （左の数列の第 1 bit）
添字|0 2 3 5 6|1 4 7|
値域|0~1|4~5|2~3|6~7|
数列|011| 5 |2 2|7 6| 01110010 （左の数列の第 0 bit）
添字|035| 1 |2 6|4 7|
値域|0|4|2 |6|1 |5|3|7|
数列|0| |22|6|11|5| |7|
添字|0| |26|7|35|1| |4|
wavelet tree の時と同じように値の範囲を分割していますが、対応する値域は値順に並んではいません。ただし、各段は安定ソートなので、各ブロックの中では添字順に並んでいるという状況は同じです。なお、最終的な数列 $(0,2,2,6,1,1,5,7)$ は、最初の$A$をビット逆順でソートしたものです。
rank 関数の役割
復習すると、wavelet tree で区間 $[L,R)$ に含まれる要素の中で$X$ 以下の個数を数えるクエリに答える場合、値域 $[0,X]$ を分割して、それぞれの管理する昇順の添字列について $[L,R)$ に入るものを数えるのでした。そして fractional cascading により索引を張ることで二分探索をなくすのでした。
wavelet matrix では、ビットベクトルの rank 関数が fractional cascading の役割を果たします！つまり、ある段の半開区間 $[L,R)$ に対応する次の段の区間の位置は次のように簡単に求まります。
- 左側の子 (bit が 0): 半開区間
[rank0(L),rank0(R))に対応 - 右側の子 (bit が 1): 半開区間
[C+rank1(L),C+rank1(R))に対応（$C$ は 0 の総数）
結果として、各段につき長さ $N$ の ビットベクトルと $0$ の総数 ($C$) だけを保持することで、rank 関数を使って添字区間を次の段に移すことができます。よって wavelet tree で対応可能なクエリを、より少ないメモリで高速に処理できます。
Wavelet matrix にデータ構造を載せる
例えば「$L,R,X$ が与えられるので、$A_L, \dots, A_{R-1}$ のうち値が $X$ 以下となる要素の総和を答えよ」という問題には、
- wavelet matrix の構築時に、要素の累積和の列を持っておく
- クエリ解答時には、
- クエリの値域を分割する
- wavelet matrix により、対応する区間を高速で計算する
- 累積和を使って要素の区間和を高速に計算する
とすれば良いです。この場合、累積和により $NB$ 個の整数を保持するので、空間計算量に注意する必要があります。
結局 wavelet matrix は対応する区間を高速で計算する機構なので、クエリに応じて、Fenwick 木やセグメント木で列集約の情報を管理することもできます。