dsalgo 0.3.10

A package for Datastructures and Algorithms.
Documentation
pub trait Monoid {
    type T;

    fn op(
        &self,
        l: Self::T,
        r: Self::T,
    ) -> Self::T;

    fn e(&self) -> Self::T;
}

pub struct Segtree<G: Monoid> {
    g: G,
    size: usize,
    data: Vec<G::T>,
}

impl<G: Monoid> Segtree<G> {
    fn n(&self) -> usize {
        self.data.len() >> 1
    }
}

use std::ops::*;

impl<G: Monoid> Index<usize> for Segtree<G> {
    type Output = G::T;

    fn index(
        &self,
        i: usize,
    ) -> &Self::Output {
        &self.data[i + self.n()]
    }
}

impl<G: Monoid> Segtree<G>
where
    G::T: Clone,
{
    pub fn size(&self) -> usize {
        self.size
    }

    fn merge(
        &mut self,
        i: usize,
    ) {
        self.data[i] =
            self.g.op(self.data[i << 1].clone(), self.data[i << 1 | 1].clone());
    }

    pub fn new(
        g: G,
        size: usize,
    ) -> Self {
        assert!(size > 0);

        let n = size.next_power_of_two();

        let data = vec![g.e(); n << 1];

        Self { g, size, data }
    }

    pub fn set(
        &mut self,
        i: usize,
        x: G::T,
    ) {
        assert!(i < self.size());

        self._set(i, 0, self.n(), 1, x);
    }

    fn _set(
        &mut self,
        i: usize,
        cl: usize,
        cr: usize,
        ci: usize,
        x: G::T,
    ) {
        assert!(cl <= i && i < cr);

        if cr - cl == 1 {
            self.data[ci] = x;

            return;
        }

        let c = (cl + cr) >> 1;

        if i < c {
            self._set(i, cl, c, ci << 1, x);
        } else {
            self._set(i, c, cr, ci << 1 | 1, x);
        }

        self.merge(ci);
    }

    pub fn fold(
        &mut self,
        l: usize,
        r: usize,
    ) -> G::T {
        assert!(l <= r && r <= self.size);

        self._fold(l, r, 0, self.n(), 1)
    }

    fn _fold(
        &mut self,
        l: usize,
        r: usize,
        cl: usize,
        cr: usize,
        i: usize,
    ) -> G::T {
        if cr <= l || r <= cl {
            return self.g.e();
        }

        if l <= cl && cr <= r {
            return self.data[i].clone();
        }

        let c = (cl + cr) >> 1;

        let vl = self._fold(l, r, cl, c, i << 1);

        let vr = self._fold(l, r, c, cr, i << 1 | 1);

        self.g.op(vl, vr)
    }
}

#[cfg(test)]

mod tests {

    use super::*;

    #[test]

    fn test() {
        struct G;

        impl Monoid for G {
            type T = i64;

            fn e(&self) -> Self::T {
                0
            }

            fn op(
                &self,
                l: Self::T,
                r: Self::T,
            ) -> Self::T {
                l + r
            }
        }

        let n = 5;

        let mut seg = Segtree::new(G {}, n);

        seg.set(2, 1);

        assert_eq!(seg[2], 1);

        assert_eq!(seg.fold(0, 5), 1);
    }
}