cp_library_rs/data_structure/
bit_2d.rs

1//! 2次元BIT
2
3use crate::algebraic_structure::monoid::Monoid;
4
5macro_rules! cfor {
6    ($def:stmt ; $fin:expr ; $incr:stmt ;; $bl:block) => {{
7        $def
8        while $fin {
9            $bl
10            $incr
11        }
12    }}
13}
14
15pub struct BIT2D<M: Monoid> {
16    pub H: usize,
17    pub W: usize,
18    pub data: Vec<Vec<M::Val>>,
19}
20
21impl<M: Monoid> BIT2D<M> {
22    #[inline]
23    fn lsb(x: usize) -> usize {
24        x & x.wrapping_neg()
25    }
26
27    /// 2次元BITを作成する
28    pub fn new(H: usize, W: usize) -> Self {
29        Self {
30            H,
31            W,
32            data: vec![vec![M::e(); W + 1]; H + 1],
33        }
34    }
35
36    /// 位置 (r,c) に値 `v` を加算する
37    /// - `(r, c)`: 加算を行うインデックス(`0-indexed`)
38    /// - `x`: 加算する値
39    pub fn add(&mut self, mut r: usize, mut c: usize, v: M::Val) {
40        // 0-indexedに修正
41        r += 1;
42        c += 1;
43
44        cfor! {let mut i = r; i <= self.H; i += Self::lsb(i) ;; {
45            cfor! {let mut j = c; j <= self.W; j += Self::lsb(j) ;; {
46                self.data[i][j] = M::op(&self.data[i][j], &v);
47            }}
48        }}
49    }
50
51    /// 左上からの和を求める
52    /// - `(r,c)`: 領域 `0<=i<r, 0<=j<c` に対しての総和(`0-indexed`)
53    pub fn prefix_sum(&self, r: usize, c: usize) -> M::Val {
54        let mut res = M::e();
55
56        cfor! {let mut i = r; i > 0; i -= Self::lsb(i) ;; {
57            cfor! {let mut j = c; j > 0; j -= Self::lsb(j) ;; {
58                res = M::op(&res, &self.data[i][j]);
59            }}
60        }}
61
62        res
63    }
64}