-
Notifications
You must be signed in to change notification settings - Fork 0
/
segment_tree.rs
105 lines (86 loc) · 2.58 KB
/
segment_tree.rs
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
use crate::monoid::Monoid;
pub struct SegmentTree<M: Monoid> {
size: usize,
tree: Vec<M::S>,
}
impl<M: Monoid<S = S>, S: Clone + Copy> SegmentTree<M> {
/// self = [e; size]
pub fn new(size: usize) -> Self {
Self {
size,
tree: vec![M::E; size << 1],
}
}
/// self = array
pub fn from(array: &[M::S]) -> Self {
let size = array.len();
let mut tree = vec![M::E; size << 1];
for i in 0..size {
tree[i + size] = array[i];
}
for i in (1..size).rev() {
tree[i] = M::op(&tree[i << 1], &tree[i << 1 | 1]);
}
return Self { size, tree };
}
/// self[i] <- s
pub fn update(&mut self, mut i: usize, s: S) {
i += self.size;
self.tree[i] = s;
while i > 1 {
i >>= 1;
self.tree[i] = M::op(&self.tree[i << 1], &self.tree[i << 1 | 1]);
}
}
/// get self[i]
pub fn get(&self, i: usize) -> S {
return self.tree[i + self.size];
}
/// calculate Π_{i ∈ range} self[i]
pub fn prod<R: std::ops::RangeBounds<usize>>(&self, range: R) -> S {
let left = match range.start_bound() {
std::ops::Bound::Included(&l) => l,
std::ops::Bound::Excluded(&l) => l + 1,
std::ops::Bound::Unbounded => 0,
};
let right = match range.end_bound() {
std::ops::Bound::Included(&r) => r + 1,
std::ops::Bound::Excluded(&r) => r,
std::ops::Bound::Unbounded => self.size,
};
return self._prod(left, right);
}
fn _prod(&self, mut left: usize, mut right: usize) -> S {
left += self.size;
right += self.size;
let (mut sl, mut sr) = (M::E, M::E);
while left < right {
if left & 1 == 1 {
sl = M::op(&sl, &self.tree[left]);
left += 1;
}
if right & 1 == 1 {
right -= 1;
sr = M::op(&self.tree[right], &sr);
}
left >>= 1;
right >>= 1;
}
return M::op(&sl, &sr);
}
}
impl<M: Monoid<S = S>, S: std::fmt::Display + Copy> std::fmt::Display for SegmentTree<M> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"{}",
(0..self.size)
.map(|i| self.get(i))
.collect::<Vec<_>>()
.iter()
.map(|x| x.to_string())
.collect::<Vec<_>>()
.join(" ")
)
}
}