競プロ用ライブラリを作る Advent Calendar 2024の22日目です.
WaveletMatrixって?
整数列に対してだいたいなんでもできるデータ構造です
整数列$x_1,x_2,...x_N$についてこういうクエリが対数時間で処理できます.
- 指定された$i$について, $x_i$を求める
- 指定された$a, b, l, r$について, 指定された区間の$x_l, x_{l+1},..., x_r$かつ$a\le x_i\le b$の範囲の数の個数を数える
- 指定された$l, r, a$について, $x_l, x_{l+1},..., x_r$の中で$a$番目に小さい/大きいものを求める
- 指定された$a, n$について, $x_i=a$を満たす小さい方/大きい方から$n$番目の$i$を求める
仕組み
完備辞書
「整数列に対してなんでもできるデータ構造」の前に「01列に対してある程度できるデータ構造」を作ります.
01列に対して以下が出来ます.
- 指定された場所が0か1か求める
- 指定された$i$について, $i$番目の0の位置を求める
- 指定された$i$について, $i$番目の1の位置を求める
- 指定された$i$について, 1番目から$i$番目までの1の個数を求める
まず, bitsetで01列を持ちます. このbitsetを64bit整数の列として表現して, 同じく64個ごとに「それより前に1がいくつあるか」を別に持っておきます.
struct SampleDict {
// bitset本体
main: Vec<u64>,
// 64個ごとに個数を持つ配列
sub: Vec<usize>,
// 1の合計の個数
total: usize,
}
impl SampleDict {
fn new(bitset: Vec<u64>) -> Self {
let mut count = 0;
let mut sub = vec![0];
for val in &bitset {
count += val.count_ones() as usize;
sub.push(count);
}
Self {
main: bitset,
sub,
total: count,
}
}
// 要素数を返す
fn len(&self) -> usize {
self.main.len() * 64
}
// index番目の値が0か1か見る
fn get(&self, index: usize) -> bool {
(self.main[index / 64] >> (index % 64)) & 1 == 1
}
// 先頭index個に1がいくつあるか数える
fn rank(&self, index: usize) -> usize {
// 64個単位の個数
let a = self.sub[index / 64] as usize;
// ↑で数え漏らした分の個数
let b = (self.main[index / 64] & !(!0 << (index % 64))).count_ones() as usize;
a + b
}
// count番目の1がどこにあるか探す
// rankを二分探索して実装します
fn select(&self, count: usize) -> usize {
let mut min = 0;
let mut max = self.main.len() * 64;
while max - min > 1 {
let mid = (min + max) / 2;
if self.rank(mid) <= count {
min = mid;
} else {
max = mid;
}
}
min
}
}
頑張ればメモリ消費量を削減したりselectのオーダーを改善したりできるらしいですが, 競プロの範囲ではそんなに問題にならないのでこのままでも大丈夫です.
本体
32bitの非負整数を持つ場合を考えます.
上位bitから順番に「bitが1になっている要素を右に寄せる」という操作を行います.
3 1 4 1 5 9 2 6 5 3 5
3 1 4 1 5 2 6 5 3 5|9 <-- 8の位が1のものを右に寄せる
3 1 1 2 3 9|4 5 6 5 5 <-- 4の位が1のものを右に寄せる
1 1 9 4 5 5 5|3 2 3 6 <-- 2の位が1のものを右に寄せる
4 2 6|1 1 9 5 5 5 3 3 <-- 1の位が1のものを右に寄せる
寄せる前の数列で注目しているbitが0か1かを書きます
0 0 0 0 0 1 0 0 0 0 0
0 0 1 0 1 0 1 1 0 1|0
1 0 0 1 1 0|0 0 1 0 0
1 1 1 0 1 1 1|1 0 1 0
これを1列ずつ完備辞書に乗せたものがWaveletMatrixです.
ここから数字を復元してみます, この01の情報から, どっちに寄せられたか, が分かるのでこれを追います.
ここの数字を復元する
v
0 0 0 0 0 1 0 0 0 0 0
| 0なので, 左に寄せた位置を見る
0 0 1 0 1 0 1 1 0 1|0
^-----v 1なので, 右に寄せた位置を見る
1 0 0 1 1 0|0 0 1 0 0
v-----^ 0なので, 左に寄せた位置を見る
1 1 1 0 1 1 1|1 0 1 0
^
01を拾っていくと0101 (10進数で5)
左に寄せた位置を見る場合は「その前にある1の個数だけ前を見る」で, 右に寄せた位置を見る場合は「その後ろにある0の個数だけ後ろを見る」とすればいいです.
fn access(dicts: &[SampleDict; 64], mut index: usize) -> u64 {
let mut ret = 0;
for (i, dict) in dicts.iter().enumerate() {
if dict.get(i) {
ret |= 1 << (63 - i);
// その前にある1の個数を消す
index -= dict.rank(index);
} else {
// 後ろにある0の個数を頑張って数えるとこうなります
index = dict.len() - dict.total - (index - dict.rank(dict.len()));
}
}
ret
}
こういう感じでいろんな機能を付けることが出来ます.
何でもできるのですが, 1つ1つ挙げていくとキリがないので省略します. 知りたい方へはこの記事とかが詳しかったです.
実装
いろいろしました. まだできることは全然あると思います.
まとめ
いかがでしたか?
明日は自明をします.