lazy_segment_tree.rs

pub mod lazy_segment_tree {
    type Range = std::ops::Range<usize>;

    pub struct LazySegmentTree<S, Op, E, F, Mapping, Composition, Id> {
        n: usize,
        size: usize,
        log: usize,
        data: Vec<S>,
        lazy: Vec<F>,
        op: Op,
        e: E,
        mapping: Mapping,
        composition: Composition,
        id: Id,
    }

    impl<S, Op, E, F, Mapping, Composition, Id> LazySegmentTree<S, Op, E, F, Mapping, Composition, Id>
    where
        S: Clone,
        E: Fn() -> S,
        F: Clone,
        Op: Fn(&S, &S) -> S,
        Mapping: Fn(&F, &S) -> S,
        Composition: Fn(&F, &F) -> F,
        Id: Fn() -> F,
    {
        pub fn new(
            n: usize,
            e: E,
            op: Op,
            mapping: Mapping,
            composition: Composition,
            id: Id,
        ) -> Self {
            let size = n.next_power_of_two() as usize;
            LazySegmentTree {
                n,
                size,
                log: size.trailing_zeros() as usize,
                data: vec![e(); 2 * size],
                lazy: vec![id(); size],
                e,
                op,
                mapping,
                composition,
                id,
            }
        }
        pub fn set(&mut self, mut index: usize, value: S) {
            assert!(index < self.n);
            index += self.size;
            for i in (1..=self.log).rev() {
                self.push(index >> i);
            }
            self.data[index] = value;
            for i in 1..=self.log {
                self.update(index >> i);
            }
        }

        pub fn get(&mut self, mut index: usize) -> S {
            assert!(index < self.n);
            index += self.size;
            for i in (1..=self.log).rev() {
                self.push(index >> i);
            }
            self.data[index].clone()
        }

        pub fn prod(&mut self, range: Range) -> S {
            let mut l = range.start;
            let mut r = range.end;
            assert!(l < r && r <= self.n);

            l += self.size;
            r += self.size;

            for i in (1..=self.log).rev() {
                if ((l >> i) << i) != l {
                    self.push(l >> i);
                }
                if ((r >> i) << i) != r {
                    self.push(r >> i);
                }
            }

            let mut sum_l = (self.e)();
            let mut sum_r = (self.e)();
            while l < r {
                if l & 1 != 0 {
                    sum_l = (self.op)(&sum_l, &self.data[l]);
                    l += 1;
                }
                if r & 1 != 0 {
                    r -= 1;
                    sum_r = (self.op)(&self.data[r], &sum_r);
                }
                l >>= 1;
                r >>= 1;
            }

            (self.op)(&sum_l, &sum_r)
        }

        pub fn all_prod(&self) -> S {
            self.data[1].clone()
        }

        pub fn apply(&mut self, mut index: usize, f: F) {
            assert!(index < self.n);
            index += self.size;
            for i in (1..=self.log).rev() {
                self.push(index >> i);
            }
            self.data[index] = (self.mapping)(&f, &self.data[index]);
            for i in 1..=self.log {
                self.update(index >> i);
            }
        }
        pub fn apply_range(&mut self, range: Range, f: F) {
            let mut l = range.start;
            let mut r = range.end;
            assert!(l <= r && r <= self.n);
            if l == r {
                return;
            }

            l += self.size;
            r += self.size;

            for i in (1..=self.log).rev() {
                if ((l >> i) << i) != l {
                    self.push(l >> i);
                }
                if ((r >> i) << i) != r {
                    self.push((r - 1) >> i);
                }
            }

            {
                let mut l = l;
                let mut r = r;
                while l < r {
                    if l & 1 != 0 {
                        self.all_apply(l, f.clone());
                        l += 1;
                    }
                    if r & 1 != 0 {
                        r -= 1;
                        self.all_apply(r, f.clone());
                    }
                    l >>= 1;
                    r >>= 1;
                }
            }

            for i in 1..=self.log {
                if ((l >> i) << i) != l {
                    self.update(l >> i);
                }
                if ((r >> i) << i) != r {
                    self.update((r - 1) >> i);
                }
            }
        }

        fn update(&mut self, k: usize) {
            self.data[k] = (self.op)(&self.data[2 * k], &self.data[2 * k + 1]);
        }
        fn all_apply(&mut self, k: usize, f: F) {
            self.data[k] = (self.mapping)(&f, &self.data[k]);
            if k < self.size {
                self.lazy[k] = (self.composition)(&f, &self.lazy[k]);
            }
        }
        fn push(&mut self, k: usize) {
            self.all_apply(2 * k, self.lazy[k].clone());
            self.all_apply(2 * k + 1, self.lazy[k].clone());
            self.lazy[k] = (self.id)();
        }
    }
}

#[cfg(test)]
mod test {
    use super::lazy_segment_tree::*;
    use rand::prelude::*;

    const INF: i64 = 1 << 60;

    #[test]
    fn edge_case() {
        let n = 5;
        let mut seg_min = LazySegmentTree::new(
            n,
            || INF,
            |&s, &t| s.min(t),
            |&f, &x| f + x,
            |&f, &g| f + g,
            || 0,
        );
        let mut values = vec![0; n];
        for i in 0..n {
            values[i] = i as i64;
            seg_min.set(i, i as i64);
        }

        let from = 1;
        let to = 4;
        let add = 2;
        for i in from..to {
            values[i] += add;
        }
        seg_min.apply_range(from..to, add);

        let pos = 2;
        let value = 1;
        let cur = seg_min.prod(pos..(pos + 1));
        seg_min.set(pos, cur - value);
        values[pos] -= value;

        for l in 0..n {
            for r in (l + 1)..(n + 1) {
                let min1 = seg_min.prod(l..r);
                let &min2 = values[l..r].iter().min().unwrap();
                assert_eq!(min1, min2);
            }
        }
    }

    #[test]
    fn random_add() {
        let mut rng = thread_rng();
        let n = 32;
        let mut array = vec![0; n];
        let mut seg_min = LazySegmentTree::new(
            n,
            || INF,
            |&s, &t| s.min(t),
            |&f, &x| f + x,
            |&f, &g| f + g,
            || 0,
        );
        let mut seg_max = LazySegmentTree::new(
            n,
            || -INF,
            |&s, &t| s.max(t),
            |&f, &x| f + x,
            |&f, &g| f + g,
            || 0,
        );
        for i in 0..n {
            let value = rng.gen_range(-1000, 1000);
            array[i] = value;
            seg_min.set(i, value);
            seg_max.set(i, value);
        }

        for l in 0..n {
            for r in (l + 1)..n {
                let value = rng.gen_range(-1000, 1000);
                seg_min.apply_range(l..r, value);
                seg_max.apply_range(l..r, value);

                for i in l..r {
                    array[i] += value;
                }

                for l in 0..n {
                    for r in (l + 1)..n {
                        let mut min = INF;
                        let mut max = -INF;
                        for i in l..r {
                            min = std::cmp::min(min, array[i]);
                            max = std::cmp::max(max, array[i]);
                        }

                        assert_eq!(seg_min.prod(l..r), min);
                        assert_eq!(seg_max.prod(l..r), max);
                    }
                }
            }
        }
    }

    #[test]
    fn random_update() {
        let mut rng = thread_rng();
        #[derive(Clone)]
        struct Num {
            len: u32,
            value: i64,
        }
        let n = 15;
        let mut array = vec![0; n];
        let mut seg_update = LazySegmentTree::new(
            n,
            || Num { len: 0, value: 0 },
            |left: &Num, right: &Num| Num {
                value: left.value * 10i64.pow(right.len) + right.value,
                len: left.len + right.len,
            },
            |f: &Option<i64>, x: &Num| {
                if let &Some(f) = f {
                    let mut value = 0;
                    for _ in 0..x.len {
                        value = value * 10 + f;
                    }
                    Num { len: x.len, value }
                } else {
                    Num {
                        len: x.len,
                        value: x.value,
                    }
                }
            },
            |f: &Option<i64>, g: &Option<i64>| {
                if f.is_some() {
                    f.clone()
                } else {
                    g.clone()
                }
            },
            || None,
        );
        for i in 0..n {
            array[i] = 1;
            seg_update.set(i, Num { len: 1, value: 1 });
        }

        for _ in 0..1000 {
            let digit = rng.gen_range(0, 10);
            let left = rng.gen_range(0, n);
            let right = rng.gen_range(left + 1, n + 1);
            for i in left..right {
                array[i] = digit;
            }
            seg_update.apply_range(left..right, Some(digit));

            let mut sum = 0;
            for i in 0..n {
                sum = sum * 10 + array[i];
            }

            assert_eq!(sum, seg_update.all_prod().value);
        }
    }
}