haar_lib/ds/
fenwick_on_fenwick.rs

1//! Fenwick木上にFenwick木を構築する。
2use crate::algo::bsearch_slice::BinarySearch;
3use std::ops::{Add, Range, RangeTo, Sub};
4
5/// [`FenwickOnFenwick`]を構築するための構造体。
6#[derive(Clone, Default)]
7pub struct FenwickOnFenwickBuilder {
8    xs: Vec<i64>,
9    ys: Vec<i64>,
10}
11
12/// Fenwick木上にFenwick木を構築する。
13#[derive(Clone)]
14pub struct FenwickOnFenwick<T> {
15    c_xs: Vec<i64>,
16    c_ys: Vec<Vec<i64>>,
17    segs: Vec<Vec<T>>,
18    zero: T,
19}
20
21impl FenwickOnFenwickBuilder {
22    /// 空の[`FenwickOnFenwickBuilder`]を用意する。
23    pub fn new() -> Self {
24        Self {
25            xs: vec![],
26            ys: vec![],
27        }
28    }
29
30    /// 点`(x, y)`を登録する。
31    pub fn add(&mut self, x: i64, y: i64) {
32        self.xs.push(x);
33        self.ys.push(y);
34    }
35
36    /// [`FenwickOnFenwick`]を構築する。
37    pub fn build<T: Copy>(self, zero: T) -> FenwickOnFenwick<T> {
38        let n = self.xs.len();
39        let mut c_xs = self.xs.clone();
40        c_xs.sort_unstable();
41        c_xs.dedup();
42
43        let x_size = c_xs.len();
44
45        let mut c_ys = vec![vec![]; x_size + 1];
46
47        let mut ord = (0..n).collect::<Vec<usize>>();
48        ord.sort_by(|&i, &j| self.ys[i].cmp(&self.ys[j]));
49
50        for i in ord {
51            let mut x = c_xs.binary_search(&self.xs[i]).unwrap() + 1;
52
53            while x <= x_size {
54                c_ys[x].push(self.ys[i]);
55                x += x & (!x + 1);
56            }
57        }
58
59        let mut segs = vec![vec![]];
60
61        for ys in c_ys.iter_mut().take(x_size + 1).skip(1) {
62            ys.dedup();
63            segs.push(vec![zero; ys.len() + 1]);
64        }
65
66        FenwickOnFenwick {
67            c_xs,
68            c_ys,
69            segs,
70            zero,
71        }
72    }
73}
74
75impl<T: Copy + Add<Output = T> + Sub<Output = T>> FenwickOnFenwick<T> {
76    /// **Time Complexity** $O(\log ^ 2 n)$
77    pub fn update(&mut self, x: i64, y: i64, value: T) {
78        let mut x = self.c_xs.binary_search(&x).unwrap() + 1;
79
80        while x <= self.c_xs.len() {
81            let mut y = self.c_ys[x].binary_search(&y).unwrap() + 1;
82
83            while y < self.segs[x].len() {
84                self.segs[x][y] = self.segs[x][y] + value;
85                y += y & (!y + 1);
86            }
87            x += x & (!x + 1);
88        }
89    }
90
91    /// **Time Complexity** $O(\log ^ 2 n)$
92    pub fn fold_2d(
93        &self,
94        Range { start: x1, end: x2 }: Range<i64>,
95        Range { start: y1, end: y2 }: Range<i64>,
96    ) -> T {
97        self.fold_to_2d(..x2, ..y2) - self.fold_to_2d(..x1, ..y2) - self.fold_to_2d(..x2, ..y1)
98            + self.fold_to_2d(..x1, ..y1)
99    }
100
101    /// **Time Complexity** $O(\log ^ 2 n)$
102    pub fn fold_to_2d(
103        &self,
104        RangeTo { end: x }: RangeTo<i64>,
105        RangeTo { end: y }: RangeTo<i64>,
106    ) -> T {
107        let mut ret = self.zero;
108        let mut x = self.c_xs.lower_bound(&x);
109
110        while x > 0 {
111            let mut y = self.c_ys[x].lower_bound(&y);
112            let seg = &self.segs[x];
113
114            while y > 0 {
115                ret = ret + seg[y];
116                y -= y & (!y + 1);
117            }
118            x -= x & (!x + 1);
119        }
120
121        ret
122    }
123}
haar_lib/ds/fenwick_on_fenwick.rs

haar_lib/ds/
fenwick_on_fenwick.rs
