use crate::monoid::Monoid;
use cargo_snippet::snippet;
#[snippet(name = "segment_tree")]
#[derive(Debug, Clone)]
pub struct SegmentTree<T: Monoid> {
len: usize,
size: usize,
segment: Vec<T>,
}
#[snippet("segment_tree")]
fn childrens_idx(n: usize) -> (usize, usize) {
(n * 2 + 1, n * 2 + 2)
}
#[snippet("segment_tree")]
fn parent_idx(n: usize) -> usize {
(n - 1) / 2
}
#[snippet("segment_tree")]
impl<T: Monoid + Clone + Copy> SegmentTree<T> {
pub fn new<I: Into<T> + Copy>(v: &[I]) -> Self {
let n = v.len();
let leaves = n.next_power_of_two();
let size = 2 * leaves - 1;
let mut value = vec![T::identity(); size];
for i in (0..size).rev() {
if i >= leaves - 1 {
if i + 1 - leaves < n {
value[i] = (v[i + 1 - leaves]).into();
} else {
continue;
}
} else {
let (left, right) = childrens_idx(i);
value[i] = T::op(&value[left], &value[right]);
}
}
Self {
len: leaves,
size,
segment: value,
}
}
pub fn len(&self) -> usize {
self.len
}
fn childrens(&self, n: usize) -> (T, T) {
let (left, right) = childrens_idx(n);
(self.segment[left], self.segment[right])
}
pub fn get(&self, i: usize) -> Option<&T> {
self.segment.get(self.size - self.len + i)
}
pub unsafe fn get_mut(&mut self, i: usize) -> Option<&mut T> {
self.segment.get_mut(self.size - i - 1)
}
pub fn update(&mut self, i: usize, v: T) {
let mut cur = self.len - self.len + i;
self.segment[cur] = v;
loop {
if cur == 0 {
break;
}
cur = parent_idx(cur);
let (left, right) = self.childrens(cur);
self.segment[cur] = T::op(&left, &right);
}
}
pub fn range(&self, from: usize, to: usize) -> T {
self.range_inner(from, to, 0, self.len, 0)
}
fn range_inner(&self, from: usize, to: usize, l_bound: usize, r_bound: usize, k: usize) -> T {
if from <= l_bound && to >= r_bound {
self.segment[k]
} else if from >= r_bound || to <= l_bound {
T::identity()
} else {
let sep = (l_bound + r_bound) / 2;
T::op(
&self.range_inner(from, to, l_bound, sep, 2 * k + 1),
&self.range_inner(from, to, sep, r_bound, 2 * k + 2),
)
}
}
}