21日目: 64AryTreeを実装する

競プロ用ライブラリを作る Advent Calendar 2024の21日目です.

64AryTreeって?

$0$以上$N$未満の整数が入る集合として

• 指定した要素の追加
• 指定した要素の削除
• 指定した要素を含むか判定
• 指定された区間と集合が共通部分を持つか判定する
• 集合の最小値を探す

などをとても定数倍が軽い$O(\log N)$で行えるデータ構造です.

仕組み

「64」はワードサイズです. JavaScriptみたいな「32bit整数の方が扱いやすい」みたいな環境だったりした場合は臨機応変に変えましょう. 以下はワードサイズが64bitだった場合の説明をします.

SegmentTreeをよく知っている方へ: boolの最大値を求めるSegmentTreeで, 64分木にしてbitsetで持つと演算が少ない命令でできてお得, という話です

64AryTreeは集合をいくつかのbitsetで表現します. 各bitsetは64bit整数の配列を内部で持っています.

// self.0[i] は64*i以上64*(i+1)未満の整数の情報を持ちます
struct SampleBitSet(Vec<u64>);
impl SampleBitSet {
    fn add(&mut self, index: usize) {
        self.0[index >> 6] |= 1 << (index & 63);
    }
    fn delete(&mut self, index: usize) {
        self.0[index >> 6] &= !(1 << (index & 63));
    }
    fn has(&self, index: usize) -> bool {
        (self.0[index >> 6] >> (index & 63)) & 1 == 1
    }
}

64AryTreeが持つ1つのbitsetはそのまま集合を表現します.

2つ目のbitsetは「1つ目のbitsetの各64bit整数の情報」を持ちます.
どういう情報を載せるかによって処理できるクエリを変えられますが, とりあえずここでは「各64bit整数が0かどうか」を載せることにします. 要素を含んでいたら対応するビットが1になるのでこれは「64bit整数が管理する値に集合に含まれるものが存在するか」を見ていることになります.

3つ目のbitsetは「2つ目のbitsetの各64bit整数の情報」を持ちます. 2つ目のときと同様です.
4つ目のbitsetは「3つ目のbitsetの各64bit整数の情報」を持ちます.
... こんな感じで持っていきます. 1つ次のbitsetの大きさはその前のbitsetの64分の1になるので, ワードサイズに収まるくらいの大きさになるところで止めます.あ

struct Sample64AryTree(Vec<SampleBitSet>);
impl Sample64AryTree {
    fn add(&mut self, mut index: usize) {
        for bitset in self.0.iter_mut() {
            bitset.add(index);
            // bitsetで操作したのは index>>6 個めの整数です
            // 次のbitsetにも「値が追加された」という情報を入れに行きます
            index >>= 6;
        }
    }
    fn delete(&mut self, mut index: usize) {
        for bitset in self.0.iter_mut() {
            bitset.add(index);
            // 操作した整数が0になりきっていない可能性があるので,
            // そういうときはbreakしてやります.
            if bitset.0[index >> 6] != 0 {
                break;
            }
            index >>= 6;
        }
    }
    fn has(&self, index: usize) -> bool {
        // 一番最初のbitsetは普通に情報を持っているので, それを見ます
        self.0[0].has(index)
    }
}

最小値を探す処理を実装しましょう.
最後の方(=要素数の少ない方)から順にbitsetを見ていきます. 最後のbitsetはワードサイズに収まるくらい小さく, 64bit整数1つだけで表現されています.
この整数の「1になってる一番下のビットの位置」を調べると, それは「1つ前のbitsetで最小の要素を持つ64bit整数の位置」になります. Rustではu64::trailing_zerosが使えます.
その「1つ前のbitsetの最小の要素を持つ64bit整数」で同じことをしてその前のbitsetの位置を調べて... ということを続ければ, 最小値が分かります

impl Sample64AryTree {
    fn min(&self) {
        let mut index = 0;
        for bitset in self.0.iter().rev() {
            index = index * 64 + (bitset.0[index].trailing_zeros() as usize);
        }
    }
}

実装

• wodsizetree.rs

このデータ構造の利点は定数倍が軽いことなので, これが生かせるようにunsafeまみれのコードを書きました.

眺めてると区間更新だとかMEXだとか個数の偶奇だとかも処理できそうな気がしてきますが, それはまた今度

まとめ

いかがでしたか?

明日は数列に何でも聞けるデータ構造を実装します.