seg_tree/segment_tree/
lazy_recursive.rs

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

seg_tree/segment_tree/
lazy_recursive.rs
