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

haar_lib/ds/
segtree_on_segtree.rs
