1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176
use crate::nodes::Node;
/// Segment tree with range queries and point updates.
/// It uses `O(n)` space, assuming that each node uses `O(1)` space.
pub struct SegmentTree<T: Node + Clone> {
nodes: Vec<T>,
n: usize,
}
impl<T> SegmentTree<T>
where
T: Node + Clone,
{
/// Builds segment tree from slice, each element of the slice will correspond to a leaf of the segment tree.
/// It has time complexity of `O(n*log(n))`, assuming that [combine](Node::combine) has constant time complexity.
pub fn build(values: &[T]) -> Self {
let n = values.len();
let mut nodes = Vec::with_capacity(2 * n);
for _ in 0..2 {
for v in values {
nodes.push(v.clone());
}
}
(1..n).rev().for_each(|i| {
nodes[i] = Node::combine(&nodes[2 * i], &nodes[2 * i + 1]);
});
Self { nodes, n }
}
/// Sets the i-th element of the segment tree to value T and update the segment tree correspondingly.
/// It will panic if i is not in `[0,n)`
/// It has time complexity of `O(log(n))`, assuming that [combine](Node::combine) has constant time complexity.
pub fn update(&mut self, i: usize, value: <T as Node>::Value) {
let mut i = i;
i += self.n;
self.nodes[i] = Node::initialize(&value);
while i > 0 {
i >>= 1;
self.nodes[i] = Node::combine(&self.nodes[2 * i], &self.nodes[2 * i + 1]);
}
}
/// Returns the result from the range `[left,right]`.
/// It returns None if and only if range is empty.
/// It will **panic** if left or right are not in `[0,n)`.
/// It has time complexity of `O(log(n))`, assuming that [combine](Node::combine) has constant time complexity.
pub fn query(&self, l: usize, r: usize) -> Option<T> {
let (mut l, mut r) = (l, r);
let mut ans_left = None;
let mut ans_right = None;
l += self.n;
r += self.n + 1;
while l < r {
if l & 1 != 0 {
ans_left = Some(match ans_left {
None => Node::initialize(self.nodes[l].value()),
Some(node) => Node::combine(&node, &self.nodes[l]),
});
l += 1;
}
if r & 1 != 0 {
r -= 1;
ans_right = Some(match ans_right {
None => Node::initialize(self.nodes[r].value()),
Some(node) => Node::combine(&self.nodes[r], &node),
});
}
l >>= 1;
r >>= 1;
}
match (ans_left, ans_right) {
(Some(ans_left), Some(ans_right)) => Some(Node::combine(&ans_left, &ans_right)),
(Some(ans_left), None) => Some(Node::initialize(ans_left.value())),
(None, Some(ans_right)) => Some(Node::initialize(ans_right.value())),
(None, None) => None,
}
}
// / A method that finds the smallest prefix[^note] `u` such that `predicate(u.value(), value)` is `true`. The following must be true:
// / - `predicate` is monotonic over prefixes[^note2].
// / - `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))`.
// /
// / These are two examples, the first is finding the smallest prefix which sums at least some value.
// / ```
// / # use seg_tree::{segment_tree::SegmentTree,default::Sum,nodes::Node};
// / let predicate = |left_value:&usize, value:&usize|{*left_value>=*value}; // Is the sum greater or equal to value?
// / let g = |left_node:&usize,value:usize|{value-*left_node}; // Subtract the sum of the prefix.
// / # let nodes: Vec<Sum<usize>> = (0..10).map(|x| Sum::initialize(&x)).collect();
// / let seg_tree = SegmentTree::build(&nodes); // [0,1,2,3,4,5,6,7,8,9] with Sum<usize> nodes
// / let index = seg_tree.lower_bound(predicate, g, 3); // Will return 2 as sum([0,1,2])>=3
// / # let sums = vec![0,1,3,6,10,15,21,28,36,45];
// / # for i in 0..10{
// / # assert_eq!(seg_tree.lower_bound(predicate, g, sums[i]), i);
// / # }
// / ```
// / The second is finding the position of the smallest value greater or equal to some value.
// / ```
// / # use seg_tree::{segment_tree::SegmentTree,default::Max,nodes::Node};
// / let predicate = |left_value:&usize, value:&usize|{*left_value>=*value}; // Is the maximum greater or equal to value?
// / let g = |_left_node:&usize,value:usize|{value}; // Do nothing
// / # let nodes: Vec<Max<usize>> = (0..10).map(|x| Max::initialize(&x)).collect();
// / let seg_tree = SegmentTree::build(&nodes); // [0,1,2,3,4,5,6,7,8,9] with Max<usize> nodes
// / let index = seg_tree.lower_bound(predicate, g, 3); // Will return 3 as 3>=3
// / # for i in 0..10{
// / # assert_eq!(seg_tree.lower_bound(predicate, g, i), i);
// / # }
// / ```
// /
// / [^note]: A prefix is a segment of the form `[0,i]`.
// /
// / [^note2]: Given two prefixes `u` and `v` if `u` is contained in `v` then `predicate(u.value(), value)` implies `predicate(v.value(), value)`.
fn lower_bound(
&self,
predicate: fn(&<T as Node>::Value, &<T as Node>::Value) -> bool,
g: fn(&<T as Node>::Value, <T as Node>::Value) -> <T as Node>::Value,
value: <T as Node>::Value,
) -> usize {
self.lower_bound_helper(1, 0, self.n - 1, predicate, g, value)
}
fn lower_bound_helper(
&self,
curr_node: usize,
i: usize,
j: usize,
predicate: fn(&<T as Node>::Value, &<T as Node>::Value) -> bool,
g: fn(&<T as Node>::Value, <T as Node>::Value) -> <T as Node>::Value,
value: <T as Node>::Value,
) -> usize {
if i == j {
return i;
}
let mid = (i + j) / 2;
let left_node = 2 * curr_node + 1;
let right_node = 2 * curr_node;
let left_value = self.nodes[left_node].value();
if predicate(left_value, &value) {
self.lower_bound_helper(left_node, i, mid, predicate, g, value)
} else {
let value = g(left_value, value);
self.lower_bound_helper(right_node, mid + 1, j, predicate, g, value)
}
}
}
#[cfg(test)]
mod tests {
use crate::{default::Min, nodes::Node};
use super::SegmentTree;
#[test]
fn non_empty_query_returns_some() {
let nodes: Vec<Min<usize>> = (0..=10).map(|x| Min::initialize(&x)).collect();
let segment_tree = SegmentTree::build(&nodes);
assert!(segment_tree.query(0, 10).is_some());
}
#[test]
fn empty_query_returns_none() {
let nodes: Vec<Min<usize>> = (0..=10).map(|x| Min::initialize(&x)).collect();
let segment_tree = SegmentTree::build(&nodes);
assert!(segment_tree.query(10, 0).is_none());
}
#[test]
fn update_works() {
let nodes: Vec<Min<usize>> = (0..=10).map(|x| Min::initialize(&x)).collect();
let mut segment_tree = SegmentTree::build(&nodes);
let value = 20;
segment_tree.update(0, value);
assert_eq!(segment_tree.query(0, 0).unwrap().value(), &value);
}
#[test]
fn query_works() {
let nodes: Vec<Min<usize>> = (0..=10).map(|x| Min::initialize(&x)).collect();
let segment_tree = SegmentTree::build(&nodes);
assert_eq!(segment_tree.query(1, 10).unwrap().value(), &1);
}
}