Lazy Segement Tree

lazy_segment_tree.rs

use std::ops::{Bound, RangeBounds};

pub struct LazySegmentTree<S, F, E, G, I> {
    table_size: usize,
    value: Vec<S>,
    op: F,
    element: E,
    lazy: Vec<Option<S>>,
    mapping: G,
    composite: I,
}

impl<S, F, E, G, I> LazySegmentTree<S, F, E, G, I>
where
    S: Copy + Eq,
    F: Fn(S, S) -> S,
    E: Fn() -> S,
    G: Fn(Option<S>, S) -> S,
    I: Fn(Option<S>, Option<S>) -> Option<S>,
{
    pub fn new(table_size: usize, table_initial: S, op: F, element: E, mapping: G, composite: I) -> Self {
        let tree_size = table_size.next_power_of_two() * 2 - 1;
        let value = vec![table_initial; tree_size];
        let lazy = vec![None; tree_size];
        Self {
            table_size,
            value,
            op,
            element,
            lazy,
            mapping,
            composite,
        }
    }

    fn get_children(&self, tree_index: usize) -> (usize, usize) {
        (tree_index * 2 + 1, tree_index * 2 + 2)
    }

    fn get_range<R>(&self, range: R) -> (usize, usize)
    where
        R: RangeBounds<usize>
    {
        let left = match range.start_bound() {
            Bound::Included(l) => *l,
            Bound::Excluded(l) => l + 1,
            Bound::Unbounded => 0,    
        };
        let right = match range.end_bound() {
            Bound::Included(r) => r + 1,
            Bound::Excluded(r) => *r,
            Bound::Unbounded => self.table_size + 1,
        };
        (left, right)
    }

    pub fn prod<R>(&mut self, range: R) -> S
    where
        R: RangeBounds<usize>
    {
        let (left, right) = self.get_range(range);
        if right == left {
            (self.element)()
        } else {
            self._prod(0, left, right, 0, self.value.len() / 2 + 1)
        }
    }

    fn _prod(
        &mut self,
        tree_index: usize,
        search_left: usize,
        search_right: usize,
        left: usize,
        right: usize,
    ) -> S {
        if search_left <= left && right <= search_right {
            self.value[tree_index]
        } else if right <= search_left || search_right <= left {
            (self.element)()
        } else {
            if self.lazy[tree_index] != None {
                self.propagate(tree_index, left, right, left, right);
            }
            let mid = (left + right) / 2;
            let (left_t_index, right_t_index) = self.get_children(tree_index);
            let l_value = self._prod(left_t_index, search_left, search_right, left, mid);
            let r_value = self._prod(right_t_index, search_left, search_right, mid, right);
            (self.op)(l_value, r_value)
        }
    }

        
    pub fn apply<R>(&mut self, v: Option<S>, range: R)
    where
        R: RangeBounds<usize>
    {
        let (left, right) = self.get_range(range);
        self._apply(v, 0, left, right, 0, self.value.len() / 2 + 1);
    }

    fn _apply(
        &mut self,
        v: Option<S>,
        tree_index: usize,
        search_left: usize,
        search_right: usize,
        left: usize,
        right: usize,
    ) {
        if right <= search_left || search_right <= left {
            return;
        }
        if search_left <= left && right <= search_right {
            self.value[tree_index] = (self.mapping)(v, self.value[tree_index]);
            self.lazy[tree_index] = (self.composite)(v, self.lazy[tree_index]);
        } else {
            if self.lazy[tree_index] != None {
                self.propagate(tree_index, left, right, left, right);
            }
            let mid = (left + right) / 2;
            let (left_t_index, right_t_index) = self.get_children(tree_index);
            self._apply(v, left_t_index, search_left, search_right, left, mid);
            self._apply(v, right_t_index, search_left, search_right, mid, right);
            self.value[tree_index] = (self.op)(self.value[left_t_index], self.value[right_t_index]);
        }
    }

    fn propagate(
        &mut self,
        tree_index: usize,
        search_left: usize,
        search_right: usize,
        left: usize,
        right: usize,
    ) {
        let lazy = self.lazy[tree_index];
        self.lazy[tree_index] = None;
        let mid = (left + right) / 2;
        let (left_t_index, right_t_index) = self.get_children(tree_index);
        self._apply(lazy, left_t_index, search_left, search_right, left, mid);
        self._apply(lazy, right_t_index, search_left, search_right, mid, right);
    }
}