haar_lib/ds/
segtree_on_segtree.rs

1//! セグメント木上にセグメント木を構築する。
2use crate::{
3    algebra::traits::*,
4    algo::{bsearch_slice::BinarySearch, merge::merge},
5    ds::segtree::*,
6};
7
8use std::ops::Range;
9
10/// [`SegtreeOnSegtree`]を構築するための構造体。
11#[derive(Clone, Default)]
12pub struct SegtreeOnSegtreeBuilder {
13    xs: Vec<i64>,
14    ys: Vec<i64>,
15}
16
17/// セグメント木上にセグメント木を構築する。
18pub struct SegtreeOnSegtree<M: Monoid> {
19    c_xs: Vec<i64>,
20    c_ys: Vec<Vec<i64>>,
21    x_size: usize,
22    segs: Vec<Option<Segtree<M>>>,
23}
24
25impl SegtreeOnSegtreeBuilder {
26    /// 空の[`SegtreeOnSegtreeBuilder`]を用意する。
27    pub fn new() -> Self {
28        Self {
29            xs: vec![],
30            ys: vec![],
31        }
32    }
33
34    /// 点`(x, y)`を登録する。
35    pub fn add(&mut self, x: i64, y: i64) {
36        self.xs.push(x);
37        self.ys.push(y);
38    }
39
40    /// [`SegtreeOnSegtree`]を構築する。
41    pub fn build<M: Monoid + Clone>(self) -> SegtreeOnSegtree<M> {
42        let n = self.xs.len();
43        let mut c_xs = self.xs.clone();
44        c_xs.sort_unstable();
45        c_xs.dedup();
46
47        let x_size = c_xs.len().next_power_of_two() * 2;
48
49        let mut c_ys = vec![vec![]; x_size];
50
51        for i in 0..n {
52            let j = c_xs.lower_bound(&self.xs[i]);
53            c_ys[j + x_size / 2].push(self.ys[i]);
54        }
55
56        for i in 0..x_size / 2 {
57            let v = &mut c_ys[i + x_size / 2];
58            v.sort();
59            v.dedup();
60        }
61
62        for i in (1..x_size / 2).rev() {
63            c_ys[i] = merge(c_ys[i << 1].clone(), c_ys[(i << 1) | 1].clone());
64            c_ys[i].dedup();
65        }
66
67        let mut segs = vec![None; x_size];
68        for i in 1..x_size {
69            segs[i] = Some(Segtree::new(c_ys[i].len()));
70        }
71
72        SegtreeOnSegtree {
73            c_xs,
74            c_ys,
75            x_size,
76            segs,
77        }
78    }
79}
80
81impl<M: Monoid + Clone> SegtreeOnSegtree<M> {
82    /// 点`(x, y)`の値を`value`で更新する。
83    pub fn update(&mut self, x: i64, y: i64, value: M) {
84        let mut i = self.c_xs.lower_bound(&x) + self.x_size / 2;
85        while i >= 1 {
86            let j = self.c_ys[i].lower_bound(&y);
87            self.segs[i].as_mut().unwrap().update(j, value.clone());
88            i >>= 1;
89        }
90    }
91
92    fn fold_sub(&self, i: usize, y1: i64, y2: i64) -> M {
93        let l = self.c_ys[i].lower_bound(&y1);
94        let r = self.c_ys[i].lower_bound(&y2);
95        self.segs[i].as_ref().unwrap().fold(l..r)
96    }
97
98    /// `[x1, x2), [y1, y2)`で計算を集約する。
99    pub fn fold_2d(
100        &self,
101        Range { start: x1, end: x2 }: Range<i64>,
102        Range { start: y1, end: y2 }: Range<i64>,
103    ) -> M {
104        let mut l = self.c_xs.lower_bound(&x1) + self.x_size / 2;
105        let mut r = self.c_xs.lower_bound(&x2) + self.x_size / 2;
106
107        let mut ret = M::id();
108
109        while l < r {
110            if r & 1 == 1 {
111                r -= 1;
112                ret = M::op(ret, self.fold_sub(r, y1, y2));
113            }
114            if l & 1 == 1 {
115                ret = M::op(ret, self.fold_sub(l, y1, y2));
116                l += 1;
117            }
118
119            l >>= 1;
120            r >>= 1;
121        }
122
123        ret
124    }
125}
haar_lib/ds/segtree_on_segtree.rs

haar_lib/ds/
segtree_on_segtree.rs
