use std::iter::FromIterator;
use crate::monoid::Monoid;
pub struct SegmentTree<M: Monoid<Id>, Id> {
pub(crate) size: usize,
pub(crate) data: Vec<M::S>,
}
impl<M, Id> std::iter::FromIterator<M::S> for SegmentTree<M, Id>
where
M: Monoid<Id>,
M::S: Clone,
{
fn from_iter<T: IntoIterator<Item = M::S>>(iter: T) -> Self {
let mut data = iter.into_iter().collect::<Vec<_>>();
let size = data.len();
let n = size.next_power_of_two();
data = (0..n)
.map(|_| M::identity())
.chain(data.into_iter())
.chain((0..n - size).map(|_| M::identity()))
.collect::<Vec<_>>();
let mut seg = Self { size, data };
(1..n).rev().for_each(|i| seg.update(i));
seg
}
}
impl<M: Monoid<Id>, Id> SegmentTree<M, Id> {
pub fn size(&self) -> usize { self.size }
pub(crate) fn n(&self) -> usize { self.data.len() >> 1 }
}
impl<M, Id> SegmentTree<M, Id>
where
M: Monoid<Id>,
M::S: Clone,
{
pub fn new<F>(size: usize, default: F) -> Self
where
F: Fn() -> M::S,
{
Self::from_iter((0..size).map(|_| default()))
}
fn update(&mut self, i: usize) {
self.data[i] = M::operate(
self.data[i << 1].clone(),
self.data[i << 1 | 1].clone(),
);
}
pub fn set(&mut self, mut i: usize, x: M::S) {
assert!(i < self.size);
i += self.n();
self.data[i] = x;
while i > 1 {
i >>= 1;
self.update(i);
}
}
pub fn reduce(&self, mut l: usize, mut r: usize) -> M::S {
assert!(l < r && r <= self.size);
let n = self.n();
l += n;
r += n;
let mut vl = M::identity();
let mut vr = M::identity();
while l < r {
if l & 1 == 1 {
vl = M::operate(vl, self.data[l].clone());
l += 1;
}
if r & 1 == 1 {
r -= 1;
vr = M::operate(self.data[r].clone(), vr);
}
l >>= 1;
r >>= 1;
}
M::operate(vl, vr)
}
}
#[cfg(test)]
mod tests {
#[test]
fn test_as_monoid() {
use crate::{
associative_property::AssociativeProperty,
binary_operation::BinaryOperation,
group_theory_id::Additive,
identity_element::IdentityElement,
};
struct Mon;
impl BinaryOperation<Additive> for Mon {
type Codomain = usize;
type Lhs = usize;
type Rhs = usize;
fn map(x: usize, y: usize) -> usize { x + y }
}
impl AssociativeProperty<Additive> for Mon {}
impl IdentityElement<Additive> for Mon {
type X = usize;
fn identity() -> usize { 0 }
}
let mut seg = super::SegmentTree::<Mon, _>::new(10, || 0);
assert_eq!(seg.reduce(0, 10), 0);
seg.set(5, 5);
assert_eq!(seg.reduce(0, 10), 5);
seg.set(5, 10);
assert_eq!(seg.reduce(0, 10), 10);
}
}