競プロのためのウェーブレット行列

ヘッダー画像はどちらかというと領域木ですね。

この記事は新歓ブログリレー2026 8日目の記事です。
前提知識：セグメントツリー
簡潔データ構造の話はしません。

Point Add Rectangle Sum

$x, y \in [0, L)$ の二次元平面上に、相異なる整数座標の点が $N$ 個あります。ここに 2 種類のクエリが $Q$ 個飛んできます。

add(i, x): $i$ 番目の点に $x$ を加算

sum(lx, ly, rx, ry): 矩形 $[l x, r x) \times [l y, ry)$ 内に含まれる点に加算された値の総和を求める

参考（少し違うけど）: https://judge.yosupo.jp/problem/point_add_rectangle_sum

この問題をウェーブレット行列を使って $O (N lo g L + Q lo g L lo g N)$ で解くことを目標にしましょう。

実は、 $⌈ lo g_{2} L ⌉$ 個の順列をうまく作ると、任意の矩形領域に含まれる点群はこれらの順列上のたかだか $2 ⌈ lo g_{2} L ⌉$ 個の区間に分割することができます。
分割さえ求められれば一点加算と区間の総和はセグ木で処理できるため、 $⌈ lo g_{2} L ⌉$ 個のセグ木を持つことで、1 クエリあたり $O (lo g L lo g N)$ で処理できます。

$⌈ lo g_{2} L ⌉$ 個の順列はどのように作ればいいでしょうか？とりあえず、 $i$ 番目の順列を次のように作ることにします。

y を幅 $2^{i}$ の短冊に分ける
各短冊の中で x 昇順に並べる
短冊ごとの列を y 昇順に連結する

（ちなみに y 昇順に連結する必要性はありません。ここの任意性が後の（定数倍）高速化に効いてきます。）

strip_count_1_2_4_combined

任意の矩形領域はたかだか $2 ⌈ lo g_{2} L ⌉$ 個の短冊に分割されます。

strip_query_split

短冊中で矩形が対応する区間は二分探索で求められるので、素朴に $O (lo g L)$ 回二分探索をすることで、分割を $O (lo g L lo g N)$ で求められます。

実は二分探索は 1 回やるだけでいいらしい...

やりたいことを整理しましょう。

wm_matrix_view

灰色部分が x 座標が $[l x, r x)$ に入っている点を表していて、 y 座標が $[l y, ry)$ のになる領域を分割するときに出てくる短冊を表すのが赤い枠線です。

親の区間 $[l, r)$ が分かっているとき、子の区間を $O (1)$ で求める方法を考えます。

rank0(i): 区間 $[0, i)$ のうち、左側に分類される要素数
rank1(i): 区間 $[0, i)$ のうち、右側に分類される要素数（ $r ank 1 (i) = i - r ank 0 (i)$ ）
z: 左側に分類される総要素数

このとき、親の区間 $[l, r)$ に対応する子の区間は次の式で求められます。

左の子の区間: $[r ank 0 (l), r ank 0 (r))$
右の子の区間: $[z + r ank 1 (l), z + r ank 1 (r))$

rank の計算は 64 個ごとに前計算して、端数部分だけ popcount で数えることで空間計算量を抑え定数倍高速化をしたいです。

そのために、短冊ごとの列を単に y 昇順で連結するのではなく、次の対応が順列全体で成り立つように順列を作ります。

順列 $i + 1$ の区間 $[l, r)$ に対して、

左側に対応する順列 $i$ の区間は $[r ank 0 (l), r ank 0 (r))$
右側に対応する順列 $i$ の区間は $[z + r ank 1 (l), z + r ank 1 (r))$

この性質を満たすには、順列 $i + 1$ から順列 $i$ を作るときに、要素を「左側に分類されるものを前、右側に分類されるものを後ろ」に安定に並べればいいです。

stable_partition_mapping

まとめると、実装したいのは次の 3 つになります。

size(): ウェーブレット行列のサイズを返す。 $O (1)$
rectangle_iter(lx, ly, rx, ry): 矩形 $[l x, r x) \times [l y, ry)$ に含まれる点を、各順列上の区間に分解してイテレータで返す。返り値の各要素 (i, l, r) は「 $i$ 番目の順列の区間 $[l, r)$ 」を表す。 $O (lo g L + lo g N)$
item_iter(x, y): 点 (x, y) が各順列でどの位置にあるかをイテレータで返す。返り値の各要素(i, j) は「 $i$ 番目の順列での位置が $j$ 」を表す。 $O (lo g L + lo g N)$

実装しましょう。
頑張って実装したのがこちらになります。

struct WaveletMatrix {
    log: usize,
    ps: Vec<(usize, usize)>,
    wm: Vec<BitVec>,
}
impl WaveletMatrix {
    fn new(ps: &[(usize, usize)]) -> Self {
        let mut ps = ps.to_vec();
        ps.sort_unstable();

        let max = ps.iter().map(|t| t.1).max().unwrap_or(0).checked_add(2).unwrap();
        let log = (usize::BITS - max.leading_zeros()) as usize;
        let len = (ps.len() + 63) / 64;

        let mut wm = vec![
            BitVec {
                zeros: 0,
                block: vec![(0, 0); len + 1]
            };
            log
        ];
        let mut a = ps.iter().map(|t| t.1 + 1).collect::<Vec<_>>();
        a.resize(len * 64, !0);

        let mut new = vec![!0; a.len()];
        for i in (0..log).rev() {
            let mut i0 = 0;
            wm[i].zeros = a.iter().map(|&t| t >> i & 1 ^ 1).sum::<_>();
            let mut i1 = wm[i].zeros;
            let mut zeros = 0;

            for j in 0..len {
                wm[i].block[j].0 = zeros;
                let mut bit = 0;

                for (k, &a) in a[j * 64..(j + 1) * 64].iter().enumerate() {
                    let f = a as u64 >> i & 1;
                    bit |= f << k;

                    let i = if f == 0 { &mut i0 } else { &mut i1 };
                    new[*i] = a;
                    *i += 1;
                }

                wm[i].block[j].1 = bit;
                zeros += bit.count_zeros();
            }

            wm[i].block[len].0 = zeros;
            std::mem::swap(&mut a, &mut new);
        }

        Self { log, ps, wm }
    }

    fn size(&self) -> (usize, usize) {
        (self.log, self.ps.len())
    }

    // i0..i1,j0..j1の矩形領域の点群を表す区間に分割
    fn rectangle_iter<RI, RJ>(
        &self,
        i: RI,
        j: RJ,
    ) -> impl Iterator<Item = (usize, std::ops::Range<usize>)>
    where
        RI: std::ops::RangeBounds<usize>,
        RJ: std::ops::RangeBounds<usize>,
    {
        fn slice_range<R>(n: usize, i: R) -> (usize, usize)
        where
            R: std::ops::RangeBounds<usize>,
        {
            use std::ops::Bound::*;
            let s = match i.start_bound() {
                Included(&a) => a,
                Excluded(&a) => a + 1,
                Unbounded => 0,
            };
            let t = match i.end_bound() {
                Included(&a) => a + 1,
                Excluded(&a) => a,
                Unbounded => n,
            };
            assert!(s <= t && t <= n);
            (s, t)
        }

        let (mut i0, mut i1) = slice_range(!0, i);
        i0 = self.ps.partition_point(|t| t.0 < i0);
        i1 = self.ps.partition_point(|t| t.0 < i1);

        let (mut j0, mut j1) = slice_range(!0, j);
        let max = (1 << self.log) - 2;
        j0 = j0.min(max);
        j1 = j1.min(max) + 1;

        let mut it = self.log - 1;
        while it != !0 && (j0 ^ j1) >> it & 1 == 0 {
            let z0 = self.wm[it].count_zeros(i0);
            let z1 = self.wm[it].count_zeros(i1);

            if j0 >> it & 1 == 0 {
                (i0, i1) = (z0, z1);
            } else {
                let med = self.wm[it].zeros;
                (i0, i1) = (med + i0 - z0, med + i1 - z1);
            }
            it -= 1;
        }

        let (mut l0, mut l1, mut r0, mut r1) = (0, 0, 0, 0);
        if it != !0 {
            let z0 = self.wm[it].count_zeros(i0);
            let z1 = self.wm[it].count_zeros(i1);
            (l0, l1) = (z0, z1);
            let med = self.wm[it].zeros;
            (r0, r1) = (med + i0 - z0, med + i1 - z1);
            it -= 1;
        }

        std::iter::from_fn(move || {
            while it != !0 {
                let z0 = self.wm[it].count_zeros(l0);
                let z1 = self.wm[it].count_zeros(l1);
                let med = self.wm[it].zeros;
                let (o0, o1) = (med + l0 - z0, med + l1 - z1);
                if j0 >> it & 1 == 0 {
                    it -= 1;
                    (l0, l1) = (z0, z1);
                    return Some((it + 1, o0..o1));
                } else {
                    it -= 1;
                    (l0, l1) = (o0, o1);
                }
            }
            None
        })
        .chain(std::iter::from_fn(move || {
            while it != !0 {
                let z0 = self.wm[it].count_zeros(r0);
                let z1 = self.wm[it].count_zeros(r1);
                let med = self.wm[it].zeros;
                if j1 >> it & 1 == 0 {
                    it -= 1;
                    (r0, r1) = (z0, z1);
                } else {
                    it -= 1;
                    (r0, r1) = (med + r0 - z0, med + r1 - z1);
                    return Some((it + 1, z0..z1));
                }
            }
            None
        }))
    }

    // 区間で点(i,j)を表す位置の集合
    fn item_iter(&self, (i, j): (usize, usize)) -> impl Iterator<Item = (usize, usize)> {
        let mut i = self.ps.binary_search(&(i, j)).unwrap();
        (0..self.log).rev().map(move |t| {
            let cnt = self.wm[t].count_zeros(i);
            if self.wm[t].get(i) {
                i = self.wm[t].zeros + i - cnt;
            } else {
                i = cnt;
            }
            (t, i)
        })
    }
}

#[derive(Clone)]
struct BitVec {
    zeros: usize,
    block: Vec<(u32, u64)>, // ここまでのzeroの個数,bit列
}
impl BitVec {
    fn get(&self, i: usize) -> bool {
        self.block[i / 64].1 >> i % 64 & 1 == 1
    }

    fn count_zeros(&self, i: usize) -> usize {
        let (cur, bit) = self.block[i / 64];
        (cur + (bit | !0 << i % 64).count_zeros()) as usize
    }
}

かなり長くなっちゃいました。実装が下手ですいません。
矩形内の個数を求めるだけだったらもっと短いんですけが、区間の分割を求めるとなるとかなり実装が嵩みそうです。
非再帰にこだわらなければもうちょっと短くすることはできるはずです。

注意として、rectangle_iter で区間を扱うとき開区間で扱いたかったので、0 もうまく扱えるようすべての要素に +1 しており、座標に usize::MAX を入れるとパニックするというのと、
オーバーフロー前提のコードを書いているので、 release で実行しないとパニックします。
あと slice_range 周りの境界処理をサボっているので rectangle_iter の引数で usize::MAX を指定してもバグると思います。

これを使ってPoint Add Rectangle Sumを解くとこうなります。

use ac_library::{Additive, Segtree};
use proconio::input;

fn main() {
    input! {
        n: usize,
        q: usize,
        ps: [(usize, usize); n],
    }

    let wm = WaveletMatrix::new(&ps);
    let (h, w) = wm.size();
    let mut segs = (0..h)
        .map(|_| Segtree::<Additive<i64>>::new(w))
        .collect::<Vec<_>>();
    let mut cur = vec![0; n];

    for _ in 0..q {
        input! {
            ty: usize,
        }

        if ty == 0 {
            input! {
                i: usize,
                x: i64,
            }

            cur[i] += x;
            for (t, pos) in wm.item_iter(ps[i]) {
                segs[t].set(pos, cur[i]);
            }
        } else {
            input! {
                lx: usize,
                ly: usize,
                rx: usize,
                ry: usize,
            }

            let mut ans = 0;
            for (t, r) in wm.rectangle_iter(lx..rx, ly..ry) {
                ans += segs[t].prod(r);
            }
            println!("{ans}");
        }
    }
}

おまけ

配列 $A$ に対する問題として見ることもできます。
$i$ 番目の要素を点 $(i, A_{i})$ に対応させればよくて、この対応を使うと、次のクエリが扱えます。

区間 $[l, r)$ で $x \leq A_{i} \leq y$ を満たす要素数
区間 $[l, r)$ の $k$ 番目の値

前者は、対応する領域を区間に分割して、その区間長の総和を取れば求まります。
後者は、ウェーブレット行列をたどりながら二分探索することで求められます（セグ木上の二分探索と似たような感じです）。

更に全要素 xor や隣接 swap なども可能ですが、整備したところで使う機会はおそらく訪れないでしょう。

明日の投稿者は@genmiraで「"なぜか授業で習わない"よくある微分方程式の解き方まとめ」です。
楽しみですね。

競プロのためのウェーブレット行列

Point Add Rectangle Sum

おまけ

この記事をシェア

関連する記事