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

haar_lib/ds/
dynamic_dual_segtree.rs
