segment_tree.rs

//! Segment tree to provide efficent range query
//!
//! Implement `Monoid` to write custom type of segment tree.
//! You can check how `RMQ` is implmented in this module.
// ref: https://github.com/asi1024/competitive-library/blob/master/cpp/include/structure/segment_tree.cpp
// ref: rust-num
use std::cmp;
use std::marker::PhantomData;

pub trait Monoid<T> {
    fn id() -> Option<T> {
        None
    }
    fn op(l: &Option<T>, r: &Option<T>) -> Option<T>;
}

pub struct MinOp<T: Ord> {
    phantom: PhantomData<T>,
}

impl<T: Ord + Clone> Monoid<T> for MinOp<T> {
    #[inline]
    fn op(l: &Option<T>, r: &Option<T>) -> Option<T> {
        match (l.clone(), r.clone()) {
            (Some(l), Some(r)) => Some(cmp::min(l, r)),
            (Some(l), None) => Some(l),
            (None, Some(r)) => Some(r),
            (None, None) => None,
        }
    }
}

pub struct SegmentTree<M: Monoid<T>, T: Clone> {
    phantom: PhantomData<M>,
    data: Vec<Option<T>>,
    size: usize,
    size_p2: usize,
}

impl<M: Monoid<T>, T: Clone> SegmentTree<M, T> {
    pub fn from_vec(v: Vec<T>) -> SegmentTree<M, T> {
        let size = v.len();
        let mut size_p2 = 1;
        while size_p2 < v.len() {
            size_p2 *= 2;
        }
        let mut data = vec![None; size_p2 * 2];
        for (i, x) in v.into_iter().enumerate() {
            data[size_p2 + i] = Some(x);
        }
        for i in (0..size_p2).rev() {
            data[i] = M::op(&data[i * 2], &data[i * 2 + 1]);
        }
        SegmentTree {
            phantom: PhantomData,
            data,
            size,
            size_p2,
        }
    }

    pub fn size(&self) -> usize {
        self.size
    }

    pub fn update(&mut self, mut pos: usize, value: T) {
        assert!(pos < self.size);
        pos += self.size_p2;
        self.data[pos] = Some(value);
        loop {
            pos /= 2;
            if pos == 0 {
                break;
            }
            self.data[pos] = M::op(&self.data[pos * 2], &self.data[pos * 2 + 1]);
        }
    }

    pub fn query(&self, mut l: usize, mut r: usize) -> Option<T> {
        assert!(l <= r && r <= self.size);
        l += self.size_p2;
        r += self.size_p2;
        let mut res1 = M::id();
        let mut res2 = M::id();
        while l < r {
            if (l & 1) != 0 {
                res1 = M::op(&res1, &self.data[l]);
                l += 1;
            }
            if (r & 1) != 0 {
                r -= 1;
                res2 = M::op(&self.data[r], &res2);
            }
            l >>= 1;
            r >>= 1;
        }
        M::op(&res1, &res2)
    }
}

/// segment tree to get minimum value in a range
pub type RMQ<T> = SegmentTree<MinOp<T>, T>;

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_rms() {
        let mut test = vec![1, 5, 4, 8, 6, 9, 2, 0, 8, 1];
        let mut rms = RMQ::from_vec(test.clone());
        for i in 0..test.len() {
            for j in i..test.len() + 1 {
                assert_eq!(test[i..j].iter().cloned().min(), rms.query(i, j));
            }
        }
        rms.update(7, 5);
        test[7] = 5;
        for i in 0..test.len() {
            for j in i..test.len() + 1 {
                assert_eq!(test[i..j].iter().cloned().min(), rms.query(i, j));
            }
        }
    }

    // Custom Monoid for sum operations
    struct SumOp<T> {
        phantom: PhantomData<T>,
    }

    impl Monoid<i32> for SumOp<i32> {
        fn id() -> Option<i32> {
            Some(0)
        }

        fn op(l: &Option<i32>, r: &Option<i32>) -> Option<i32> {
            match (l, r) {
                (Some(l), Some(r)) => Some(l + r),
                (Some(l), None) => Some(*l),
                (None, Some(r)) => Some(*r),
                (None, None) => Some(0),
            }
        }
    }

    #[test]
    fn test_sum_segment_tree() {
        let v = vec![1, 2, 3, 4, 5];
        let mut seg: SegmentTree<SumOp<i32>, i32> = SegmentTree::from_vec(v.clone());

        // Test range queries
        assert_eq!(seg.query(0, 5), Some(15)); // Sum of all elements
        assert_eq!(seg.query(0, 3), Some(6)); // 1 + 2 + 3
        assert_eq!(seg.query(2, 5), Some(12)); // 3 + 4 + 5
        assert_eq!(seg.query(1, 4), Some(9)); // 2 + 3 + 4
        assert_eq!(seg.query(2, 2), Some(0)); // Empty range

        // Test update
        seg.update(2, 10); // Change 3 to 10
        assert_eq!(seg.query(0, 5), Some(22)); // 1 + 2 + 10 + 4 + 5
        assert_eq!(seg.query(2, 3), Some(10)); // Just the updated element
    }

    // Custom Monoid for max operations
    struct MaxOp<T: Ord> {
        phantom: PhantomData<T>,
    }

    impl<T: Ord + Clone> Monoid<T> for MaxOp<T> {
        fn op(l: &Option<T>, r: &Option<T>) -> Option<T> {
            match (l.clone(), r.clone()) {
                (Some(l), Some(r)) => Some(cmp::max(l, r)),
                (Some(l), None) => Some(l),
                (None, Some(r)) => Some(r),
                (None, None) => None,
            }
        }
    }

    #[test]
    fn test_max_segment_tree() {
        let v = vec![3, 1, 4, 1, 5, 9, 2, 6];
        let mut seg: SegmentTree<MaxOp<i32>, i32> = SegmentTree::from_vec(v);

        assert_eq!(seg.query(0, 8), Some(9)); // Max of all
        assert_eq!(seg.query(0, 3), Some(4)); // Max of [3, 1, 4]
        assert_eq!(seg.query(4, 6), Some(9)); // Max of [5, 9]

        seg.update(5, 2); // Change 9 to 2
        assert_eq!(seg.query(0, 8), Some(6)); // Max is now 6
        assert_eq!(seg.query(4, 6), Some(5)); // Max of [5, 2]
    }

    // Custom Monoid for product operations with modulo
    struct ModProdOp {
        phantom: PhantomData<i64>,
    }

    impl Monoid<i64> for ModProdOp {
        fn id() -> Option<i64> {
            Some(1)
        }

        fn op(l: &Option<i64>, r: &Option<i64>) -> Option<i64> {
            const MOD: i64 = 1_000_000_007;
            match (l, r) {
                (Some(l), Some(r)) => Some((l * r) % MOD),
                (Some(l), None) => Some(*l),
                (None, Some(r)) => Some(*r),
                (None, None) => Some(1),
            }
        }
    }

    #[test]
    fn test_product_segment_tree() {
        let v = vec![2, 3, 5, 7, 11];
        let seg: SegmentTree<ModProdOp, i64> = SegmentTree::from_vec(v);

        assert_eq!(seg.query(0, 5), Some(2310)); // 2 * 3 * 5 * 7 * 11
        assert_eq!(seg.query(1, 4), Some(105)); // 3 * 5 * 7
        assert_eq!(seg.query(0, 1), Some(2)); // Just 2
        assert_eq!(seg.query(3, 3), Some(1)); // Empty range returns identity
    }

    #[test]
    fn test_size_method() {
        let v = vec![1, 2, 3, 4, 5];
        let seg: SegmentTree<MinOp<i32>, i32> = SegmentTree::from_vec(v);
        assert_eq!(seg.size(), 5);

        let empty_seg: SegmentTree<MinOp<i32>, i32> = SegmentTree::from_vec(vec![]);
        assert_eq!(empty_seg.size(), 0);
    }

    #[test]
    fn test_single_element() {
        let v = vec![42];
        let mut seg: SegmentTree<MinOp<i32>, i32> = SegmentTree::from_vec(v);
        assert_eq!(seg.query(0, 1), Some(42));

        seg.update(0, 100);
        assert_eq!(seg.query(0, 1), Some(100));
    }

    #[test]
    fn test_power_of_two_size() {
        // Test with exact power of 2 size
        let v = vec![1, 2, 3, 4, 5, 6, 7, 8];
        let seg: SegmentTree<MinOp<i32>, i32> = SegmentTree::from_vec(v);
        assert_eq!(seg.query(0, 8), Some(1));
        assert_eq!(seg.query(4, 8), Some(5));
    }

    #[test]
    #[should_panic]
    fn test_update_out_of_bounds() {
        let v = vec![1, 2, 3];
        let mut seg: SegmentTree<MinOp<i32>, i32> = SegmentTree::from_vec(v);
        seg.update(3, 4); // Index 3 is out of bounds
    }

    #[test]
    #[should_panic]
    fn test_query_invalid_range() {
        let v = vec![1, 2, 3];
        let seg: SegmentTree<MinOp<i32>, i32> = SegmentTree::from_vec(v);
        seg.query(2, 1); // l > r
    }

    #[test]
    #[should_panic]
    fn test_query_out_of_bounds() {
        let v = vec![1, 2, 3];
        let seg: SegmentTree<MinOp<i32>, i32> = SegmentTree::from_vec(v);
        seg.query(0, 4); // r > size
    }

    #[test]
    fn test_empty_vector() {
        let seg: SegmentTree<MinOp<i32>, i32> = SegmentTree::from_vec(vec![]);
        assert_eq!(seg.size(), 0);
        assert_eq!(seg.query(0, 0), None);
    }

    // Test with custom struct
    #[derive(Clone, Ord, PartialOrd, Eq, PartialEq, Debug)]
    struct Point {
        x: i32,
        y: i32,
    }

    #[test]
    fn test_custom_struct() {
        let points = vec![
            Point { x: 1, y: 2 },
            Point { x: 3, y: 1 },
            Point { x: 2, y: 3 },
        ];
        let mut seg: SegmentTree<MinOp<Point>, Point> = SegmentTree::from_vec(points);

        assert_eq!(seg.query(0, 3), Some(Point { x: 1, y: 2 }));

        seg.update(0, Point { x: 5, y: 0 });
        assert_eq!(seg.query(0, 3), Some(Point { x: 2, y: 3 }));
    }
}