まえがき
この記事は投稿者(NokonoKotlin)の個人サイトの記事から Qiita 用に移植 & 加筆したものです。 (この文言は記事が剽窃でないことを周知するためのものであり、個人サイトの宣伝ではないことをあらかじめご了承ください。)
はじめに
Wavelet Matrix なんて全然怖くないのに... 天下り的に手続きを読むと理解に苦しむだけで、はじめにちゃんと Wavelet Matrix のアイデアを理解さえすれば、Wavelet Matrix 上の操作は全て直感で理解できるはずなんでち。
目次
-
ビットベクトル
- 概要
- 空間計算量
-
Wavelet Matrix
- アイデア
- 構築
- Wavelet Matrix
-
二次元矩形取得
- 座標圧縮
- Wavelet Matrix にのせる
1 ビットベクトル
$01$ 列のことをビットベクトルと呼びます。適切に実装すれば、長さ $N$ のビットベクトルのデータサイズは $N$ ビットです。ここでは、長さ $N$ のビットベクトルに対して以下の操作を定義します。
- $\mathrm{rank0}(B,i) := $ ビットベクトル $B$ の先頭 $i$ 個のビットに含まれる $0$ の個数
- $\mathrm{rank1}(B,i) := $ ビットベクトル $B$ の先頭 $i$ 個のビットに含まれる $1$ の個数
1.1 概要
この処理を行うために以下の図のようなデータ構造を作ります。
まず、$\mathrm{lb}$ ビットごとにそれまでの $1$ のビットの個数を記録した配列 $L$ を作ります。この配列のサイズは $\frac{N}{lb}$ です。
次に、$\mathrm{b}$ ビットごとに「最後に $L$ に記録された位置」以降の $1$ のビットの個数を記録した配列 $S$ を作ります。この配列のサイズは $\frac{N}{b}$ です。ただし、$L$ と重なる部分に記録する値は未定義とします。実装しやすいように自身でよしなにしてください。
最後に、元のビット列を $\mathrm{b}$ ビットごと $\mathrm{b}$ ビット整数にビットマスクしておきます。ただし、図では左のビットが小さい桁になるようにビットマスクしています。
このデータ構造を用いることで、先述の $\mathrm{rank1}$ の計算を以下のように実行できます。
配列 $L$ と配列 $S$ を使用可能な部分に関しては、それらから $1$ のビットの個数を取得します。その後、はみ出ている部分 (長さを $x$ とする) の $1$ の個数を、ビットマスクした整数の先頭 $x$ ビットの $\mathrm{pop\_count}$ で取得し、それらを全て足し合わせます。
$\mathrm{pop\_count}$ は整数の $1$ のビットの個数を取得する操作なので、先頭 $x$ ビットに対してこれを実行するためにビットシフトを行う必要があります。厳密にはわかりませんが、ビットシフトも $\mathrm{pop\_count}$ も高速に実行可能です。
1.2 空間計算量
配列 $L$ のサイズは $O(\frac{N}{lb})$ であり、要素の大きさは $O(N)$ なので $O(\frac{N\log{N}}{lb})$ ビット必要です。また配列 $S$ も同様に $O(\log{lb})$ ビットを $O(\frac{N}{b})$ 個格納するので $O(\frac{N\log{lb}}{b})$ ビット必要です。
$b,lb$ はデータ構造のパラメータなので、これを適切に設定することで空間計算量を $O(N)$ ビットに抑えることができます。自分は $lb = 256 , b = 16$ としています。
2 Wavelet Matrix
さて、ここで本記事の主役の登場です。Wavelet Matrix です。いきますよ!! いいですか!? いきますよ!!?!? ばーーーーーーん!!!!!!! ブラウザバックはもう少々お持ちください。見た目に惑わされてはいけません。Wavelet Matrix は見た目は怖いけど実は優しいタイプなのです。
2.1 アイデア
Wavelet Matrix は非負整数からなる数列 $A$ を元に構築するものであり、図では最上段の (data) が数列 $A = \{1,0,9,2,6,4,12,0,9,7,9,6\}$ に対応します。また、Wavelet Matrix は $\lceil \log{(\mathrm{max}(A))}\rceil + 1$ 段からなります。
各段に格納されているそれぞれのデータはビットのパターンによってブロックに分類されます。図で示す矢印は、その矢印の根元の要素が次の段でどのビットパターンのブロックに分類されるかを示しており、最上段からの矢印の経路がビットパターンに対応します。
破線の矢印は次のビットが $0$ のパターンの移動先を示しており、実践の矢印は $1$ のパターンの移動先を示しています。例えば以下の図では、最上段から $01$ のパターンに対応する矢印の経路を辿ることで、数列 $A$ の要素のうちビットパターンが $01??$ の要素のみを格納したブロックにたどりつきます。
最上段から矢印の経路を辿ることはビットパターンの上の桁からビットを確定していく意味を持ち、段々とパターンに当てはまる整数の取りうる値の範囲が二分されていくことから、Wavelet Matrix は実は数列要素の値の範囲を分割統治するセグメントツリーのような構造であることが理解できます。
このことがわかると、数列 $A$ の値の範囲に関する条件を付きのクエリを実行することができます。セグメント木のように、調べたい値の範囲に収まる値のみを含むブロックを探索し、見つけた $O(\log{(\mathrm{max}(A))})$ 個のブロックごとに処理を行うことができます。
また数列 $A$ の値を条件にするのとは反対に、数列 $A$ の要素の中から何らかの性質(条件)を満たす要素を見つける操作も可能です。
例えば最上段からの探索において、未到達のブロックのうち含まれる要素数が最大のものを優先して探索を行う場合を考えます(探索済み経路の末端に隣接するブロックを優先度付きキューに入れて探索)。最下段のブロックはビットパターンが全て確定しており、含まれる要素は全てそのビットパターンと一致することから、要素数が最大のブロックを優先する探索で初めて到達した最下段のブロックのビットパターンは、登場する値の最頻値となります。
例えば他にも、要素数が $k$ 以上のブロックのみを探索し、移動ルートは $0$ の矢印を優先することで、$A$ の要素のうち登場回数が $k$ 以上の要素の最小値が求まります。ただし、これらの操作はいずれも Wavelet Matrix 上での操作として可能なものの例を示すためのものであり、競技プログラミングで使用するには時間計算量が大きいという点に注意が必要です。
このように Wavelet Matrix での探索方法をうまく工夫することで、数列 $A$ から特定の要素及びそれに付随するデータを取得することが可能です。できることのバリエーションは沢山ありますが、Wavelet Matrix のアイデアさえ理解できていれば特に困ることはないと思います。
2.2 構築
まず、各段の各要素に対して、その段までで確定したビットパターンの次の桁のビットをビットベクトルで管理します。たとえば最上段のビットパターンは何も確定していないので、最上段では各要素の最上位ビットの情報をビットベクトルで管理します。同様に次の段では最上ビットの次のビットを管理します。
ある段の各ブロックは、そのブロックのビットパターンで明らかになっている最右のビットが $0$ なら左に、$1$ なら右に配置されます。要素のビットパターンを分類したものがブロックなので、一つ上の段で管理しているビットに注目して要素の安定ソートを行うことで、この並び順を得ることができます。また、安定ソートなのでビットが同じ要素同士は一つ上の段の順序を保持したまま配置されます。
ビットの $0/1$ でソートすることがブロックの分割に寄与し、先述の通り値の範囲の分割統治の構造が構築されます。また安定ソートなので同じブロック内の要素の順は元の数列 $A$ における登場順と同じになります。このことから、数列 $A$ の指定した区間内であるという条件を保持したまま Wavelet Matrix を探索することが可能になります。以下の図は $A$ の区間 $[1,8)$ 内の要素のみに注目して探索を行う様子です。
各段の要素はビットの $0/1$ で左右に分けられるという情報から、移動先のブロックにおける (元の数列 $A$ でのポジションを保持した) ポジションを復元することができます。図で示す通り、$\mathrm{BV}[p]$ を管理する段で $i$ 番目のポジションは次の段で以下のポジションに移動します。
- $0$ の矢印を辿る場合 : $\mathrm{rank0}(\mathrm{BV}[p] , i)$
- $1$ の矢印を辿る場合 : $\mathrm{rank0}(\mathrm{BV}[p] , |A|) + \mathrm{rank1}(\mathrm{BV}[p] , i)$
このように構築した Wavelet Matrix では、これまでに述べたとおり数列 $A$ に対して、指定した区間内の要素に対して、値の範囲を指定してデータを取得することができます。たとえば以下の図は $A$ の区間 $[1,8)$ 内の要素で $6$ 以下の要素に注目して Wavelet Matrix を探索する様子です。
より一般に、指定する値の範囲を $[a,b)$ とするとき、調べているブロックのビットパターンを持つ整数の範囲と $[a,b)$ を比較して探索を枝刈りすることで、探索セグメントツリーと同様の計算量で終了することがわかります。
- 調べているブロックのビットパターンを持つ整数が全て値の範囲 $[a,b)$ に収まるなら、調べている領域の要素は全て条件を満たすので採用し、このパターン以降の探索を打ち切る。
- ビットパターンの範囲と値の範囲が重ならない時、調べている領域に条件を満たす要素は存在しないので、このパターン以降の探索を打ち切る。
- 上記のいずれでもない時、すなわち今調べている領域に条件を満たす要素と満たさない要素のどちらも含まれている可能性がある時、$0$ の行き先と $1$ の行き先の二手に別れて探索を続行する。
また、前節の最後で説明したことと同様に、数列 $A$ の指定した区間内の要素の中から何らかの性質(条件)を満たす要素を見つける操作も可能です。例えば前節と同様に最頻値を求める場合、$A$ の指定した区間内の要素を最も多く含むブロックを優先して探索を行うことで、指定した区間内という条件付きで最頻値を求めることができます。各ブロックの $A$ の指定した区間に含まれる要素数は、そのブロック内の区間のポジションを先述の通りに復元して計算することができます。
2.3 Wavelet Matrix
さて皆さんお気づきでしょうか。これまで説明した処理では、値の範囲の絞り込みをビットパターンで行い、区間の条件の保持をビットベクトルの $\mathrm{rank}$ 機能で行いました。つまり、これまで説明した処理に必要なのものは実は各段のビットベクトルのみなのです。
よって Wavelet Matrix は以下のような構造になります。ただし、ビットパターンを全て確定したブロックのために、最下段に仮想的な段を設けています。
$O(|A|)$ ビットのビットベクトルが $\lceil \log{(\mathrm{max}(A))}\rceil$ 段あるので、Wavelet Matrix のビット数は $O(|A|\log{(\mathrm{max}(A))})$ ビットとなり、できることに対してかなり省メモリなデータ構造であることがわかります。
先述の通り、これらのビットベクトルの情報のみで先ほどと同様の探索を実行することができます。以下は先ほどと同様に $A$ の区間 $[1,8)$ 内の要素で $6$ 以下の要素に注目して Wavelet Matrix を探索する様子です。
また、数列 $A$ の各段における並び順を記録し、それぞれの段ごとに区間取得データ構造を用意することで、Wavelet Matrix 上の探索で辿り着いた {段 , 区間} の情報からさまざまなデータを取得することができます。
たとえば各段の並び順で累積和テーブルを用意した場合、数列 $A$ に対して区間と値の範囲の条件を満たす要素の総和を取得することができます。また、$A$ が数列ではなくペアの列の場合、$A$ の一つ目の値の範囲を指定し、二つ目の値の集約を取得するなどの処理を行うこともできます。
ただし用意するデータ構造の空間計算量に注意する必要があります。たとえば上記の累積和のテーブル $1$ つのビット数は $O(|A|\log{(\mathrm{max}(A))})$ ビットなので、全ての段についてこれを用意すると全体で $O(|A|\log^{2}{(\mathrm{max}(A))})$ ビットになります。累積和であればそこまで大きくはなりませんが、これがもし Sparse Table なら $O(|A|\log{|A|}\log^{2}{(\mathrm{max}(A))})$ ビット必要になります。
Sparse Table の代わりとして、Cartesian Tree を用いて定数時間で区間 Min/Max を取得できる省メモリなデータ構造があるらしいですが、Wavelet Matrix とは別の話なので興味があれば各自記事などを探してみてください。
3 二次元矩形取得
$xy$ 座標系上の重み付きの $N$ 個の点に対して、矩形内の点の重みの集約を取得する処理を Wavelet Matrix を用いて実現します。
3.1 座標圧縮
与えられた点の $x$ 座標と $y$ 座標をそれぞれ独立に以下を定義します。
- $\mathrm{comp}(X) := N$ 個の点のうち、$x$ 座標が $X$ 未満の点の個数。
- $\mathrm{comp\_y}(Y) := N$ 個の点のうち、$y$ 座標が $Y$ 未満の点の $y$ 座標の種類数。
今回、以下の図のように点の座標 $(X,Y)$ を $(\mathrm{comp}(X),\mathrm{comp\_y}(Y))$ に変換したものを考えます。
この方法のもとでは、元の座標上の点の $x$ 座標の条件 $l \leq x \leq r$ を、圧縮後の座標 $x\prime$ に関して $\mathrm{comp}(l) \leq x\prime \leq \mathrm{comp}(r)$ と言い換えることができます。$y$ 座標についても同様です ($\mathrm{comp\_y}$ で変換する)。
3.2 Wavelet Matrix にのせる
座標圧縮した点を $x$ 座標の昇順でソートし、その順で点の $y$ 座標を Wavelet Matrix で表現します。またこのとき、各段にはその段における要素の並び順で点の重みを保存します。
すると、座標圧縮後の $x$ 座標に関する範囲の条件と Wavelet Matrix が表現する点の並び順における区間の条件がマッチします。すなわち、元の座標で $x$ 座標が $l$ 以上 $r$ 未満の点を対象にする場合、Wavelet Matrix 上の区間 $[\mathrm{comp}(l) , \mathrm{comp}(r))$ を調べれば良いのです。
$y$ 座標の範囲の条件に関しても同様に $\mathrm{comp\_y}$ で変換し、Wavelet Matrix 上で値の条件で絞り込みながらブロックを探索します。あとは普通の Wavelet Matrix と同様に探索し、採用した {段 , 区間} の情報に対して重みの区間取得を行うことで、点の重みの集約を行うことができます。
ソースコード
ライセンスは守ってくださいね (ライブラリ冒頭や最後にライセンスのコメントを書いてあります)。ライセンスに従わない利用を確認した場合、ライセンスに従わない利用をやめるよう指示するなどの処置を取る場合があります。
#include<vector>
#include<cassert>
#include<iostream>
#include<stack>
#include<bit>
using std::vector,std::stack;
using std::pair,std::tuple,std::get;
using std::popcount;
using std::cout,std::endl,std::cin;
/*
Copyright ©️ (c) NokonoKotlin (okoteiyu) 2024.
Released under the MIT license(https://opensource.org/licenses/mit-license.php)
*/
struct BitVector{
private:
inline static const int m_B_ = 64;
using bitmask_int = uint64_t;
inline static const int m_LB_ = 65536;
using smallblock_int = uint16_t;
unsigned int m_N_;
vector<bitmask_int> m_V_;
vector<unsigned int> m_L_;
vector<smallblock_int> m_S_;
public:
BitVector(){}
BitVector(unsigned int n):m_N_(n){
m_V_.resize((m_N_+m_B_-1)/m_B_,0);
m_L_.resize((m_N_+m_LB_-1)/m_LB_ + 1,0);
m_S_.resize((m_N_+m_B_-1)/m_B_ + 1,0);
}
void build(){
unsigned int bitsum = 0;
m_L_[0] = 0;
m_S_[0] = 0;
for(unsigned int i = 0 ; i < m_N_ ; i++){
if(this->get(i))bitsum++;
if((i+1)%m_LB_ == 0)m_L_[(i+1)/m_LB_] = bitsum;
if((i+1)%m_B_ == 0)m_S_[(i+1)/m_B_] = bitsum - m_L_[(i+1)/m_LB_];
}
}
unsigned int size() const& {
return m_N_;
}
void set1(unsigned int i){
assert(0<=i&&i<m_N_);
m_V_[i/m_B_] |= (bitmask_int(1)<<(i%m_B_));
}
bool get(unsigned int i) const& {
assert(0<=i&&i<m_N_);
return bool(m_V_[i/m_B_]&(bitmask_int(1)<<(i%m_B_)));
}
bool operator[](unsigned int i) const& {
return this->get(i);
}
unsigned int rank1(unsigned int i) const& {
assert(0<=i&&i<=m_N_);
unsigned int res = m_L_[i/m_LB_] + m_S_[(i/m_B_)] ;
if(i%m_B_ != 0)res += popcount<bitmask_int>(bitmask_int(bitmask_int(m_V_[i/m_B_]) << (m_B_-i%m_B_)));
return res;
}
unsigned int rank0(unsigned int i) const& {
assert(0<=i&&i<=m_N_);
return i - rank1(i);
}
/*
Copyright ©️ (c) NokonoKotlin (okoteiyu) 2024.
Released under the MIT license(https://opensource.org/licenses/mit-license.php)
*/
};
/*
Copyright ©️ (c) NokonoKotlin (okoteiyu) 2025.
Released under the MIT license(https://opensource.org/licenses/mit-license.php)
*/
template<typename int_type = long long, int _maxbit_ = 61>
class WaveletMatrix{
private:
static_assert((int_type(1)<<(_maxbit_+1)) > 0);
int m_N;
BitVector m_BV[_maxbit_+1];
bool built = false;
public:
vector<vector<int>> build(vector<int_type> A){
m_N = A.size();
vector<vector<int>> res(_maxbit_+2,vector<int>(m_N));
for(int i = 0 ; i < m_N ; i++)res[_maxbit_+1][i] = i;
for(int b = _maxbit_ ; b >= 0 ; b--){
m_BV[b] = BitVector(m_N);
vector<pair<int_type,int>> lef,rig;
for(int i = 0 ; i < A.size() ; i++ ){
if((int_type(1)<<b)&(A[i])){
m_BV[b].set1(i);
rig.emplace_back(A[i],res[b+1][i]);
}else {
lef.emplace_back(A[i],res[b+1][i]);
}
}
for(int i = 0 ; i < lef.size() ; i++){
A[i] = lef[i].first;
res[b][i] = lef[i].second;
}
for(int i = 0 ; i < rig.size() ; i++){
A[lef.size()+i] = rig[i].first;
res[b][lef.size()+i] = rig[i].second;
}
m_BV[b].build();
}
return res;
}
vector<tuple<int,int,int>> segments(int L , int R , int_type P , int_type Q) const& {
assert(L < R);
assert(P < Q);
vector<tuple<int,int,int>> ranges;
stack<tuple<int,int,int,int_type,int_type>> st;
st.emplace(_maxbit_+1,L,R,0,(int_type(1)<<(_maxbit_+1)));
while(!st.empty()){
tuple<int,int,int,long long , long long> now = st.top();
st.pop();
int b = std::get<0>(now);
int l = std::get<1>(now);
int r = std::get<2>(now);
assert(l < r);
int_type b_ly = std::get<3>(now);
int_type b_ry = std::get<4>(now);
if(b_ry <= P || Q <= b_ly)continue;
if(P <= b_ly && b_ry <= Q){
ranges.emplace_back(b,l,r);
continue;
}
assert(1 <= b);
int l0 = m_BV[b-1].rank0(l);
int r0 = m_BV[b-1].rank0(r);
int l1 = m_BV[b-1].rank0(m_BV[b-1].size()) + m_BV[b-1].rank1(l);
int r1 = m_BV[b-1].rank0(m_BV[b-1].size()) + m_BV[b-1].rank1(r);
if(l0 < r0)st.emplace(b-1,l0,r0,b_ly,(b_ly+b_ry)>>1);
if(l1 < r1)st.emplace(b-1,l1,r1,(b_ly+b_ry)>>1,b_ry);
}
return ranges;
}
inline static constexpr int maxbit(){
return _maxbit_;
}
int size() const& {
return m_N;
}
const BitVector& operator[](int p) const& {
return m_BV[p];
}
/*
Copyright ©️ (c) NokonoKotlin (okoteiyu) 2025.
Released under the MIT license(https://opensource.org/licenses/mit-license.php)
*/
};
使い方
記事で説明した通り、ビットベクトルを並べた最低限の構造だけのライブラリです。ただし、ウェーブレット行列の構築時に各段の要素の並び順を、元の列に対する添字(0-index)の順列で返します。この順列は各段に対して計算するので、合計で $\lceil \log{(\mathrm{max}(A))}\rceil + 1$ 個の順列を返します。この順列を用いて累積和などのテーブルを構築することができます。
-
[b]
:= $b$ 桁目のビットのビットベクトルのデータ構造にアクセス -
segments(L,R,P,Q)
:= 元の列の区間 $[L,R)$ 内の要素で、$P$ 以上 $Q$ 未満の要素を全て過不足なく格納する領域のデータたちを以下の形式で返す。-
(b,l,r)
:= 段 $b$ の区間 $[l,r)$ の要素が全て条件を満たす
-
使用例 1
Library Checker の個数カウントクエリを通しました。
// verified by (https://judge.yosupo.jp/problem/static_range_frequency)
void StaticRangeFrequency(){
WaveletMatrix W;
int n,q; cin >> n >> q;
vector<long long> a(n);
for(long long& x : a)cin >> x;
// a のウェーブレット行列を構築
W.build(a);
for(int i = 0 ; i < q ; i++){
int l,r;
long long x;
cin >> l >> r >> x;
// 整数 x のビットパターンに沿って、区間 [l,r) のポジションを保持したまま探索する
for(int b = W.maxbit() ; b >= 0 ; b--){
if((1ll<<b)&x){
l = W[b].rank0(n)+W[b].rank1(l);
r = W[b].rank0(n)+W[b].rank1(r);
}else{
l = W[b].rank0(l);
r = W[b].rank0(r);
}
}
// 最下段のブロックは x に対応するので、その中での [l,r) のポジションの要素数が答え
cout << r - l << endl;
}
}
使用例 2
Library Checker の矩形取得クエリを通しました。
// verified by (https://judge.yosupo.jp/problem/rectangle_sum)
void RectangleSum(){
int n; cin >> n;
int q; cin >> q;
vector<tuple<int , int , long long>> Points;
WaveletMatrix<int,29> W;
vector<vector<long long>> prefix_sum(31,vector<long long>(n+1,0));
vector<int> X(n),Y(n);
for(int i = 0 ; i < n ; i++){
long long w;
cin >> X[i] >> Y[i] >> w;
Points.emplace_back(X[i],Y[i],w);
}
// 座標圧縮のためにソート
sort(X.begin(),X.end());
sort(Y.begin(),Y.end());
for(int i = 0 ; i < n ; i++){
auto&& [x,y,w] = Points[i];
// 点の座標を圧縮
x = int(lower_bound(X.begin(),X.end(),x) - X.begin());
// 記事で書いたものとは異なる圧縮方法であることに注意 (x の方法を採用)
y = int(lower_bound(Y.begin(),Y.end(),y) - Y.begin());
Points[i] = {x,y,w};
}
// 点を x 座標の昇順でソート
sort(Points.begin() , Points.end() ,
[&](tuple<int,int,long long> a , tuple<int,int,long long> b){
return bool(get<0>(a) < get<0>(b));
}
);
vector<int> nY;
for(int i = 0 ; i < n ; i++)nY.push_back(get<1>(Points[i]));
// 点の Y 座標の値でウェーブレット行列を構築し、各段の並び順を記録しておく
vector<vector<int>> ordered = W.build(nY);
// 各段の点の並びにおける重みの累積和テーブルを構築
for(int b = 0 ; b < ordered.size() ; b++){
for(int i = 0 ; i < n ; i++){
// ordered の順で累積和テーブルを構築
prefix_sum[b][i+1] = prefix_sum[b][i] + get<2>(Points[ordered[b][i]]);
}
}
// クエリ処理
for(int i = 0 ; i < q ; i++){
int l,d,r,u;
cin >> l >> d >> r >> u;
// 座標圧縮後の座標の条件に変換
l = int(lower_bound(X.begin(),X.end(),l) - X.begin());
r = int(lower_bound(X.begin(),X.end(),r) - X.begin());
d = int(lower_bound(Y.begin(),Y.end(),d) - Y.begin());
u = int(lower_bound(Y.begin(),Y.end(),u) - Y.begin());
// ライブラリは区間のサイズが正であることを要求するので、ここで条件分岐
if(l == r || d == u){
cout << 0 << endl;
continue;
}
// 条件を満たす要素を格納する {段番号 , 左端 , 右端} の集合を取得
vector<tuple<int,int,int>> segments = W.segments(l,r,d,u);
long long sum = 0;
// 取得した区間に対して累積和を取得
for(tuple<int,int,int> seg : segments){
auto&& [b,l,r] = seg;
sum += prefix_sum[b][r] - prefix_sum[b][l];
}
cout << sum << endl;
}
}