haar_lib/ds/
dynamic_dual_segtree.rs

1//! 動的双対セグメント木
2use crate::algebra::traits::Monoid;
3use crate::misc::nullable_usize::NullableUsize;
4use std::ops::Range;
5
6#[derive(Clone, Debug)]
7struct Node<T> {
8    value: T,
9    left: NullableUsize,
10    right: NullableUsize,
11}
12
13impl<T> Node<T> {
14    fn new(value: T) -> Self {
15        Self {
16            value,
17            left: NullableUsize::NULL,
18            right: NullableUsize::NULL,
19        }
20    }
21}
22
23/// 動的双対セグメント木
24#[derive(Clone, Debug)]
25pub struct DynamicDualSegtree<M: Monoid> {
26    data: Vec<Node<M>>,
27    root: NullableUsize,
28    to: usize,
29}
30
31impl<M: Monoid + Clone> Default for DynamicDualSegtree<M> {
32    fn default() -> Self {
33        Self::new()
34    }
35}
36
37impl<M: Monoid + Clone> DynamicDualSegtree<M> {
38    /// [`DynamicDualSegtree<M>`]を生成する。
39    pub fn new() -> Self {
40        Self {
41            data: vec![Node::new(M::id())],
42            root: NullableUsize(0),
43            to: 1,
44        }
45    }
46
47    fn propagate(&mut self, cur: usize, from: usize, to: usize) {
48        if to - from > 1 {
49            let mut cur_node = self.data[cur].clone();
50            let value = cur_node.value.clone();
51
52            let left = cur_node.left;
53            if left.is_null() {
54                let t = Node::new(M::id());
55                cur_node.left = NullableUsize(self.data.len());
56                self.data.push(t);
57            }
58            let left = cur_node.left.0;
59            let lv = self.data[left].value.clone();
60            self.data[left].value = M::op(lv, value.clone());
61
62            let right = cur_node.right;
63            if right.is_null() {
64                let t = Node::new(M::id());
65                cur_node.right = NullableUsize(self.data.len());
66                self.data.push(t);
67            }
68            let right = cur_node.right.0;
69            let rv = self.data[right].value.clone();
70            self.data[right].value = M::op(rv, value);
71
72            cur_node.value = M::id();
73            self.data[cur] = cur_node;
74        }
75    }
76
77    #[allow(clippy::collapsible_else_if)]
78    fn update_(&mut self, cur: usize, from: usize, to: usize, s: usize, t: usize, value: &M) {
79        if to - from == 1 {
80            if s <= from && to <= t {
81                let cur_value = unsafe { self.data.get_unchecked(cur).value.clone() };
82                unsafe {
83                    self.data.get_unchecked_mut(cur).value = M::op(cur_value, value.clone());
84                }
85            }
86        } else {
87            if to < s || t < from {
88            } else if s <= from && to <= t {
89                let cur_value = unsafe { self.data.get_unchecked(cur).value.clone() };
90                unsafe {
91                    self.data.get_unchecked_mut(cur).value = M::op(cur_value, value.clone());
92                }
93            } else {
94                let mid = (from + to) / 2;
95                self.propagate(cur, from, to);
96                let cur = unsafe { self.data.get_unchecked(cur) };
97                let left = cur.left;
98                let right = cur.right;
99                self.update_(left.0, from, mid, s, t, value);
100                self.update_(right.0, mid, to, s, t, value);
101            }
102        }
103    }
104
105    /// 範囲`s..t`を`value`で更新する。
106    pub fn update(&mut self, Range { start: s, end: t }: Range<usize>, value: M) {
107        loop {
108            let root = self.root.0;
109
110            if t <= self.to {
111                break;
112            }
113            self.to *= 2;
114
115            let mut new_root = Node::new(M::id());
116            new_root.left = NullableUsize(root);
117
118            self.root = NullableUsize(self.data.len());
119            self.data.push(new_root);
120        }
121
122        self.update_(self.root.0, 0, self.to, s, t, &value);
123    }
124
125    fn get_(&mut self, cur: usize, from: usize, to: usize, i: usize) -> M {
126        if !(from..to).contains(&i) {
127            return M::id();
128        }
129
130        if to - from == 1 {
131            unsafe { self.data.get_unchecked(cur).value.clone() }
132        } else {
133            self.propagate(cur, from, to);
134
135            let mid = (from + to) / 2;
136            let cur = unsafe { self.data.get_unchecked(cur) };
137            if i < mid {
138                self.get_(cur.left.0, from, mid, i)
139            } else {
140                self.get_(cur.right.0, mid, to, i)
141            }
142        }
143    }
144
145    /// `i`番目の要素を取得する。
146    pub fn get(&mut self, i: usize) -> M {
147        self.get_(self.root.0, 0, self.to, i)
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154    use crate::algebra::sum::*;
155    use my_testtools::*;
156    use rand::Rng;
157
158    #[test]
159    fn test() {
160        let n = 1000;
161
162        let mut a = vec![Sum::id(); n];
163        let mut seg = DynamicDualSegtree::<Sum<u32>>::new();
164
165        let mut rng = rand::thread_rng();
166
167        for _ in 0..100 {
168            let lr = rand_range(&mut rng, 0..n);
169            let x = rng.gen_range(0..10000);
170
171            seg.update(lr.clone(), Sum(x));
172            a[lr].iter_mut().for_each(|e| e.op_assign_r(Sum(x)));
173
174            assert_eq!(a, (0..n).map(|i| seg.get(i)).collect::<Vec<_>>());
175        }
176    }
177}
haar_lib/ds/dynamic_dual_segtree.rs

haar_lib/ds/
dynamic_dual_segtree.rs
