seg_tree/segment_tree/
iterative.rs

1use core::mem::MaybeUninit;
2
3use crate::nodes::Node;
4
5/// Segment tree with range queries and point updates.
6/// It uses `O(n)` space, assuming that each node uses `O(1)` space.
7/// Note if you need to use `lower_bound`, just use the [`RecursiveSegmentTree`](crate::segment_tree::RecursiveSegmentTree) it uses double the memory though and it's less performant.
8pub struct Iterative<T> {
9    nodes: Vec<T>,
10    n: usize,
11}
12
13impl<T> Iterative<T>
14where
15    T: Node + Clone,
16{
17    /// Builds segment tree from slice, each element of the slice will correspond to a leaf of the segment tree.
18    /// It has time complexity of `O(n*log(n))`, assuming that [combine](Node::combine) has constant time complexity.
19    pub fn build(values: &[T]) -> Self {
20        let n = values.len();
21        let mut nodes: Vec<MaybeUninit<T>> = Vec::with_capacity(2 * n);
22        unsafe { nodes.set_len(2 * n) };
23        for i in 0..n {
24            nodes[i + n].write(values[i].clone());
25        }
26        for i in (1..n).rev() {
27            let (bottom_nodes, top_nodes) = nodes.split_at_mut(i + 1);
28            bottom_nodes[i].write(Node::combine(
29                unsafe { top_nodes[i - 1].assume_init_ref() },
30                unsafe { top_nodes[i].assume_init_ref() },
31            ));
32        }
33        let ptr = nodes.as_mut_ptr();
34        core::mem::forget(nodes);
35        let nodes = unsafe { Vec::from_raw_parts(ptr.cast(), 2 * n, 2 * n) };
36        Self { nodes, n }
37    }
38
39    /// Sets the i-th element of the segment tree to value T and update the segment tree correspondingly.
40    /// It will panic if i is not in `[0,n)`
41    /// It has time complexity of `O(log(n))`, assuming that [combine](Node::combine) has constant time complexity.
42    pub fn update(&mut self, i: usize, value: &<T as Node>::Value) {
43        let mut i = i;
44        i += self.n;
45        self.nodes[i] = Node::initialize(value);
46        i >>= 1;
47        while i > 0 {
48            self.nodes[i] = Node::combine(&self.nodes[2 * i], &self.nodes[2 * i + 1]);
49            i >>= 1;
50        }
51    }
52
53    /// Returns the result from the range `[left,right]`.
54    /// It returns None if and only if range is empty.
55    /// It will **panic** if left or right are not in `[0,n)`.
56    /// It has time complexity of `O(log(n))`, assuming that [combine](Node::combine) has constant time complexity.
57    #[allow(clippy::must_use_candidate)]
58    pub fn query(&self, l: usize, r: usize) -> Option<T> {
59        let (mut l, mut r) = (l, r);
60        let mut ans_left = None;
61        let mut ans_right = None;
62        l += self.n;
63        r += self.n + 1;
64        while l < r {
65            if l & 1 != 0 {
66                ans_left = Some(match ans_left {
67                    None => Node::initialize(self.nodes[l].value()),
68                    Some(node) => Node::combine(&node, &self.nodes[l]),
69                });
70                l += 1;
71            }
72            if r & 1 != 0 {
73                r -= 1;
74                ans_right = Some(match ans_right {
75                    None => Node::initialize(self.nodes[r].value()),
76                    Some(node) => Node::combine(&self.nodes[r], &node),
77                });
78            }
79            l >>= 1;
80            r >>= 1;
81        }
82        match (ans_left, ans_right) {
83            (Some(ans_left), Some(ans_right)) => Some(Node::combine(&ans_left, &ans_right)),
84            (Some(ans_left), None) => Some(Node::initialize(ans_left.value())),
85            (None, Some(ans_right)) => Some(Node::initialize(ans_right.value())),
86            (None, None) => None,
87        }
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use crate::{nodes::Node, utils::Min};
94
95    use super::Iterative;
96
97    #[test]
98    fn non_empty_query_returns_some() {
99        let nodes: Vec<Min<usize>> = (0..=10).map(|x| Min::initialize(&x)).collect();
100        let segment_tree = Iterative::build(&nodes);
101        assert!(segment_tree.query(0, 10).is_some());
102    }
103    #[test]
104    fn empty_query_returns_none() {
105        let nodes: Vec<Min<usize>> = (0..=10).map(|x| Min::initialize(&x)).collect();
106        let segment_tree = Iterative::build(&nodes);
107        assert!(segment_tree.query(10, 0).is_none());
108    }
109    #[test]
110    fn update_works() {
111        let nodes: Vec<Min<usize>> = (0..=10).map(|x| Min::initialize(&x)).collect();
112        let mut segment_tree = Iterative::build(&nodes);
113        let value = 20;
114        segment_tree.update(0, &value);
115        assert_eq!(segment_tree.query(0, 0).unwrap().value(), &value);
116    }
117    #[test]
118    fn query_works() {
119        let nodes: Vec<Min<usize>> = (0..=10).map(|x| Min::initialize(&x)).collect();
120        let segment_tree = Iterative::build(&nodes);
121        for i in 0..10 {
122            assert_eq!(segment_tree.query(i, 10).unwrap().value(), &i);
123        }
124    }
125}
seg_tree/segment_tree/iterative.rs

seg_tree/segment_tree/
iterative.rs
