seg_tree/segment_tree/
lazy_persistent.rs

1use crate::nodes::{LazyNode, Node, PersistentNode};
2
3/// Lazy persistent segment tree, it saves every version of itself, it has range queries and range updates.
4/// It uses `O(n+q*log(n))` space, where `q` is the amount of updates, and assuming that each node uses `O(1)` space.
5pub struct LazyPersistent<T: PersistentNode + LazyNode> {
6    nodes: Vec<T>,
7    roots: Vec<usize>,
8    n: usize,
9}
10
11impl<T> LazyPersistent<T>
12where
13    T: PersistentNode + LazyNode + Clone,
14{
15    /// Builds a lazy persistent segment tree from slice, each element of the slice will correspond to a leaf of the segment tree.
16    /// It has time complexity of `O(n*log(n))`, assuming that [combine](Node::combine) has constant time complexity.
17    pub fn build(values: &[T]) -> Self {
18        let n = values.len();
19        let mut temp = Self {
20            nodes: Vec::with_capacity(4 * n),
21            roots: Vec::with_capacity(1),
22            n,
23        };
24        if n == 0 {
25            return temp;
26        }
27        let root = temp.build_helper(values, 0, n - 1);
28        temp.roots.push(root);
29        temp
30    }
31
32    fn build_helper(&mut self, values: &[T], i: usize, j: usize) -> usize {
33        let mid = (i + j) / 2;
34        if i == j {
35            let curr_node = self.nodes.len();
36            self.nodes.push(values[i].clone());
37            return curr_node;
38        }
39        let left_node = self.build_helper(values, i, mid);
40        let right_node = self.build_helper(values, mid + 1, j);
41        let curr_node = self.nodes.len();
42        self.nodes
43            .push(Node::combine(&self.nodes[left_node], &self.nodes[right_node]));
44        self.nodes[curr_node].set_children(left_node, right_node);
45        curr_node
46    }
47
48    /// Returns the result from the range `[left,right]` from the version of the segment tree.
49    /// It returns None if and only if range is empty.
50    /// It will **panic** if left or right are not in [0,n), or if version is not in [0,[versions](LazyPersistentSegmentTree::versions)).
51    /// It has time complexity of `O(log(n))`, assuming that [combine](Node::combine), [`update_lazy_value`](LazyNode::update_lazy_value) and [`lazy_update`](LazyNode::lazy_update) have constant time complexity.
52    pub fn query(&mut self, version: usize, left: usize, right: usize) -> Option<T> {
53        self.query_helper(self.roots[version], left, right, 0, self.n - 1)
54    }
55
56    fn push(&mut self, curr_node: usize, i: usize, j: usize) {
57        if self.nodes[curr_node].lazy_value().is_some() && i != j {
58            let left_node = self.nodes.len();
59            let right_node = self.nodes.len() + 1;
60            self.nodes
61                .push(self.nodes[self.nodes[curr_node].left_child()].clone());
62            self.nodes
63                .push(self.nodes[self.nodes[curr_node].right_child()].clone());
64            let (parent_slice, sons_slice) = self.nodes.split_at_mut(curr_node + 1);
65            let value = parent_slice[curr_node].lazy_value().unwrap();
66            sons_slice[left_node - curr_node - 1].update_lazy_value(value);
67            sons_slice[right_node - curr_node - 1].update_lazy_value(value);
68        }
69        self.nodes[curr_node].lazy_update(i, j);
70    }
71
72    fn query_helper(
73        &mut self,
74        curr_node: usize,
75        left: usize,
76        right: usize,
77        i: usize,
78        j: usize,
79    ) -> Option<T> {
80        if j < left || right < i {
81            return None;
82        }
83        if self.nodes[curr_node].lazy_value().is_some() {
84            self.push(curr_node, i, j);
85        }
86        if left <= i && j <= right {
87            return Some(self.nodes[curr_node].clone());
88        }
89        let mid = (i + j) / 2;
90        let left_node = self.nodes[curr_node].left_child();
91        let right_node = self.nodes[curr_node].right_child();
92        match (
93            self.query_helper(left_node, left, right, i, mid),
94            self.query_helper(right_node, left, right, mid + 1, j),
95        ) {
96            (Some(ans_left), Some(ans_right)) => Some(Node::combine(&ans_left, &ans_right)),
97            (Some(ans_left), None) => Some(ans_left),
98            (None, Some(ans_right)) => Some(ans_right),
99            (None, None) => None,
100        }
101    }
102
103    /// Creates a new segment tree version from version were the p-th element of the segment tree to value T and update the segment tree correspondingly.
104    /// It will panic if p is not in `[0,n)`, or if version is not in [0,[versions](LazyPersistentSegmentTree::versions)).
105    /// It has time complexity of `O(log(n))`, assuming that [combine](Node::combine), [`update_lazy_value`](LazyNode::update_lazy_value) and [`lazy_update`](LazyNode::lazy_update) have constant time complexity.
106    pub fn update(&mut self, version: usize, left: usize, right: usize, value: &<T as Node>::Value) {
107        let new_root = self.update_helper(self.roots[version], left, right, value, 0, self.n - 1);
108        self.roots.push(new_root);
109    }
110
111    fn update_helper(
112        &mut self,
113        curr_node: usize,
114        left: usize,
115        right: usize,
116        value: &<T as Node>::Value,
117        i: usize,
118        j: usize,
119    ) -> usize {
120        if j < left || right < i {
121            return curr_node;
122        }
123        let x = self.nodes.len();
124        self.nodes.push(self.nodes[curr_node].clone());
125        if left <= i && j <= right {
126            self.nodes[x].update_lazy_value(value);
127            self.push(x, i, j);
128            return x;
129        }
130        let mid = (i + j) / 2;
131        let left_node = self.update_helper(self.nodes[x].left_child(), left, right, value, i, mid);
132        let right_node =
133            self.update_helper(self.nodes[x].right_child(), left, right, value, mid + 1, j);
134        self.nodes[x] = Node::combine(&self.nodes[left_node], &self.nodes[right_node]);
135        self.nodes[x].set_children(left_node, right_node);
136        x
137    }
138
139    /// Return the amount of different versions the current segment tree has.
140    #[allow(clippy::must_use_candidate)]
141    pub fn versions(&self) -> usize {
142        self.roots.len()
143    }
144
145    /// A method that finds the smallest prefix[^note] `u` such that `predicate(u.value(), value)` is `true`. The following must be true:
146    /// - `predicate` is monotonic over prefixes[^note2].
147    /// - `g` will satisfy the following, given segments `[i,j]` and `[i,k]` with `j<k` we have that `predicate([i,k].value(),value)` implies `predicate([j+1,k].value(),g([i,j].value(),value))`.
148    ///
149    /// These are two examples, the first is finding the smallest prefix which sums at least some value.
150    /// ```
151    /// # use seg_tree::{LazyPersistent,utils::{Sum, PersistentWrapper},nodes::Node};
152    /// # type PSum<T> = PersistentWrapper<Sum<T>>;
153    /// let predicate = |left_value: &usize, value: &usize|{ *left_value >= *value }; // Is the sum greater or equal to value?
154    /// let g = |left_node: &usize, value: usize|{ value - *left_node }; // Subtract the sum of the prefix.
155    /// # let nodes: Vec<PSum<usize>> = (0..10).map(|x| PSum::initialize(&x)).collect();
156    /// let seg_tree = LazyPersistent::build(&nodes); // [0,1,2,3,4,5,6,7,8,9] with Sum<usize> nodes
157    /// let index = seg_tree.lower_bound(0, predicate, g, 3); // Will return 2 as sum([0,1,2])>=3
158    /// # let sums = vec![0,1,3,6,10,15,21,28,36,45];
159    /// # for i in 0..10{
160    /// #    assert_eq!(seg_tree.lower_bound(0, predicate, g, sums[i]), i);
161    /// # }
162    /// ```
163    /// The second is finding the position of the smallest value greater or equal to some value.
164    /// ```
165    /// # use seg_tree::{LazyPersistent,utils::{Max, PersistentWrapper, LazySetWrapper},nodes::Node};
166    /// # type PMax<T> = PersistentWrapper<LazySetWrapper<Max<T>>>;
167    /// let predicate = |left_value:&usize, value:&usize|{*left_value>=*value}; // Is the maximum greater or equal to value?
168    /// let g = |_left_node:&usize,value:usize|{value}; // Do nothing
169    /// # let nodes: Vec<PMax<usize>> = (0..10).map(|x| PMax::initialize(&x)).collect();
170    /// let seg_tree = LazyPersistent::build(&nodes); // [0,1,2,3,4,5,6,7,8,9] with Max<usize> nodes
171    /// let index = seg_tree.lower_bound(0, predicate, g, 3); // Will return 3 as 3>=3
172    /// # for i in 0..10{
173    /// #    assert_eq!(seg_tree.lower_bound(0, predicate, g, i), i);
174    /// # }
175    /// ```
176    ///
177    /// [^note]: A prefix is a segment of the form `[0,i]`.
178    ///
179    /// [^note2]: Given two prefixes `u` and `v` if `u` is contained in `v` then `predicate(u.value(), value)` implies `predicate(v.value(), value)`.
180    pub fn lower_bound<F, G>(
181        &self,
182        version: usize,
183        predicate: F,
184        g: G,
185        value: <T as Node>::Value,
186    ) -> usize
187    where
188        F: Fn(&<T as Node>::Value, &<T as Node>::Value) -> bool,
189        G: Fn(&<T as Node>::Value, <T as Node>::Value) -> <T as Node>::Value,
190    {
191        self.lower_bound_helper(self.roots[version], 0, self.n - 1, predicate, g, value)
192    }
193    fn lower_bound_helper<F, G>(
194        &self,
195        curr_node: usize,
196        i: usize,
197        j: usize,
198        predicate: F,
199        g: G,
200        value: <T as Node>::Value,
201    ) -> usize
202    where
203        F: Fn(&<T as Node>::Value, &<T as Node>::Value) -> bool,
204        G: Fn(&<T as Node>::Value, <T as Node>::Value) -> <T as Node>::Value,
205    {
206        if i == j {
207            return i;
208        }
209        let mid = (i + j) / 2;
210        let left_node = self.nodes[curr_node].left_child();
211        let right_node = self.nodes[curr_node].right_child();
212        let left_value = self.nodes[left_node].value();
213        if predicate(left_value, &value) {
214            self.lower_bound_helper(left_node, i, mid, predicate, g, value)
215        } else {
216            let value = g(left_value, value);
217            self.lower_bound_helper(right_node, mid + 1, j, predicate, g, value)
218        }
219    }
220}
221#[cfg(test)]
222mod tests {
223    use crate::{
224        nodes::Node,
225        segment_tree::lazy_persistent::LazyPersistent,
226        utils::{PersistentWrapper, Sum},
227    };
228    type PSum<T> = PersistentWrapper<Sum<T>>;
229    #[test]
230    fn non_empty_query_returns_some() {
231        let nodes: Vec<PSum<usize>> = (0..=10).map(|x| PSum::initialize(&x)).collect();
232        let mut segment_tree = LazyPersistent::build(&nodes);
233        assert!(segment_tree.query(0, 0, 10).is_some());
234    }
235    #[test]
236    fn empty_query_returns_none() {
237        let nodes: Vec<PSum<usize>> = (0..=10).map(|x| PSum::initialize(&x)).collect();
238        let mut segment_tree = LazyPersistent::build(&nodes);
239        assert!(segment_tree.query(0, 10, 0).is_none());
240    }
241    #[test]
242    fn normal_update_works() {
243        let nodes: Vec<PSum<usize>> = (0..=10).map(|x| PSum::initialize(&x)).collect();
244        let mut segment_tree = LazyPersistent::build(&nodes);
245        let value = 20;
246        segment_tree.update(0, 0, 0, &value);
247        assert_eq!(segment_tree.query(1, 0, 0).unwrap().value(), &value);
248    }
249
250    #[test]
251    fn branched_update_works() {
252        let nodes: Vec<PSum<usize>> = (0..=10).map(|x| PSum::initialize(&x)).collect();
253        let mut segment_tree = LazyPersistent::build(&nodes);
254        let value = 20;
255        segment_tree.update(0, 0, 10, &value);
256        segment_tree.update(0, 1, 1, &value);
257        assert_eq!(segment_tree.query(2, 0, 0).unwrap().value(), &0);
258        assert_eq!(segment_tree.query(2, 1, 1).unwrap().value(), &(value + 1));
259    }
260
261    #[test]
262    fn query_works() {
263        let nodes: Vec<PSum<usize>> = (0..=10).map(|x| PSum::initialize(&x)).collect();
264        let mut segment_tree = LazyPersistent::build(&nodes);
265        assert_eq!(segment_tree.query(0, 0, 10).unwrap().value(), &55);
266    }
267}
seg_tree/segment_tree/lazy_persistent.rs

seg_tree/segment_tree/
lazy_persistent.rs
