Implement FromIterator for Lazy/Segtree and improve performance

src/lazysegtree.rs

@@ -31,12 +31,39 @@ impl<F: MapMonoid> LazySegtree<F> {
 }
 impl<F: MapMonoid> From<Vec<<F::M as Monoid>::S>> for LazySegtree<F> {
     fn from(v: Vec<<F::M as Monoid>::S>) -> Self {
-        let n = v.len();
+        Self::from_vec(v, 0)
+    }
+}
+impl<F: MapMonoid> FromIterator<<F::M as Monoid>::S> for LazySegtree<F> {
+    fn from_iter<T: IntoIterator<Item = <F::M as Monoid>::S>>(iter: T) -> Self {
+        let iter = iter.into_iter();
+        let n = iter.size_hint().0;
+        let log = ceil_pow2(n as u32) as usize;
+        let size = 1 << log;
+        let mut d = Vec::with_capacity(size * 2);
+        d.extend(repeat_with(F::identity_element).take(size).chain(iter));
+        Self::from_vec(d, size)
+    }
+}
+
+impl<F: MapMonoid> LazySegtree<F> {
+    /// Creates a segtree from elements `d[offset..]`.
+    fn from_vec(mut d: Vec<<F::M as Monoid>::S>, offset: usize) -> Self {
+        assert!(offset <= d.len());
+        let n = d.len() - offset;
         let log = ceil_pow2(n as u32) as usize;
         let size = 1 << log;
-        let mut d = vec![F::identity_element(); 2 * size];
+        match offset.cmp(&size) {
+            Ordering::Less => {
+                d.splice(0..0, repeat_with(F::identity_element).take(size - offset));
+            }
+            Ordering::Equal => {}
+            Ordering::Greater => {
+                d.splice(size..offset, empty());
+            }
+        };
+        d.resize_with(size * 2, F::identity_element);
         let lz = vec![F::identity_map(); size];
-        d[size..(size + n)].clone_from_slice(&v);
         let mut ret = LazySegtree {
             n,
             size,
@@ -47,11 +74,11 @@ impl<F: MapMonoid> From<Vec<<F::M as Monoid>::S>> for LazySegtree<F> {
         for i in (1..size).rev() {
             ret.update(i);
         }
+        // `ret.d[0]` is uninitialized and has an unknown value.
+        // This is ok as it is unused (as of writing).
         ret
     }
-}
 
-impl<F: MapMonoid> LazySegtree<F> {
     pub fn set(&mut self, mut p: usize, x: <F::M as Monoid>::S) {
         assert!(p < self.n);
         p += self.size;
@@ -323,7 +350,9 @@ where
 
 // TODO is it useful?
 use std::{
+    cmp::Ordering,
     fmt::{Debug, Error, Formatter, Write},
+    iter::{empty, repeat_with, FromIterator},
     ops::{Bound, RangeBounds},
 };
 impl<F> Debug for LazySegtree<F>
@@ -354,7 +383,7 @@ where
 mod tests {
     use std::ops::{Bound::*, RangeBounds};
 
-    use crate::{LazySegtree, MapMonoid, Max};
+    use crate::{Additive, LazySegtree, MapMonoid, Max};
 
     struct MaxAdd;
     impl MapMonoid for MaxAdd {
@@ -410,6 +439,14 @@ mod tests {
         check_segtree(&internal, &mut segtree);
     }
 
+    #[test]
+    fn test_from_iter() {
+        let it = || (1..7).map(|x| x * 4 % 11);
+        let base = it().collect::<Vec<_>>();
+        let mut segtree: LazySegtree<MaxAdd> = it().collect();
+        check_segtree(&base, &mut segtree);
+    }
+
     //noinspection DuplicatedCode
     fn check_segtree(base: &[i32], segtree: &mut LazySegtree<MaxAdd>) {
         let n = base.len();
@@ -466,6 +503,48 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_from_vec() {
+        struct IdAdditive;
+        impl MapMonoid for IdAdditive {
+            type M = Additive<u32>;
+            type F = ();
+            fn identity_map() {}
+            fn mapping(_: &(), &x: &u32) -> u32 {
+                x
+            }
+            fn composition(_: &(), _: &()) {}
+        }
+
+        let v = vec![1, 2, 4];
+        let ans_124 = vec![7, 3, 4, 1, 2, 4, 0];
+        let tree = LazySegtree::<IdAdditive>::from_vec(v, 0);
+        assert_eq!(&tree.d[1..], &ans_124[..]);
+
+        let v = vec![1, 2, 4, 8];
+        let tree = LazySegtree::<IdAdditive>::from_vec(v, 0);
+        assert_eq!(&tree.d[1..], &vec![15, 3, 12, 1, 2, 4, 8][..]);
+
+        let v = vec![1, 2, 4, 8, 16];
+        let tree = LazySegtree::<IdAdditive>::from_vec(v, 0);
+        assert_eq!(
+            &tree.d[1..],
+            &vec![31, 15, 16, 3, 12, 16, 0, 1, 2, 4, 8, 16, 0, 0, 0][..]
+        );
+
+        let v = vec![314, 159, 265, 1, 2, 4];
+        let tree = LazySegtree::<IdAdditive>::from_vec(v, 3);
+        assert_eq!(&tree.d[1..], &ans_124[..]);
+
+        let v = vec![314, 159, 265, 897, 1, 2, 4];
+        let tree = LazySegtree::<IdAdditive>::from_vec(v, 4);
+        assert_eq!(&tree.d[1..], &ans_124[..]);
+
+        let v = vec![314, 159, 265, 897, 932, 1, 2, 4];
+        let tree = LazySegtree::<IdAdditive>::from_vec(v, 5);
+        assert_eq!(&tree.d[1..], &ans_124[..]);
+    }
+
     fn check(base: &[i32], segtree: &mut LazySegtree<MaxAdd>, range: impl RangeBounds<usize>) {
         let expected = base
             .iter()

src/segtree.rs

@@ -1,7 +1,8 @@
 use crate::internal_bit::ceil_pow2;
 use crate::internal_type_traits::{BoundedAbove, BoundedBelow, One, Zero};
-use std::cmp::{max, min};
+use std::cmp::{max, min, Ordering};
 use std::convert::Infallible;
+use std::iter::{empty, repeat_with, FromIterator};
 use std::marker::PhantomData;
 use std::ops::{Add, Bound, Mul, RangeBounds};
 
@@ -80,19 +81,46 @@ impl<M: Monoid> Segtree<M> {
 }
 impl<M: Monoid> From<Vec<M::S>> for Segtree<M> {
     fn from(v: Vec<M::S>) -> Self {
-        let n = v.len();
+        Self::from_vec(v, 0)
+    }
+}
+impl<M: Monoid> FromIterator<M::S> for Segtree<M> {
+    fn from_iter<T: IntoIterator<Item = M::S>>(iter: T) -> Self {
+        let iter = iter.into_iter();
+        let n = iter.size_hint().0;
+        let log = ceil_pow2(n as u32) as usize;
+        let size = 1 << log;
+        let mut d = Vec::with_capacity(size * 2);
+        d.extend(repeat_with(M::identity).take(size).chain(iter));
+        Self::from_vec(d, size)
+    }
+}
+impl<M: Monoid> Segtree<M> {
+    /// Creates a segtree from elements `d[offset..]`.
+    fn from_vec(mut d: Vec<M::S>, offset: usize) -> Self {
+        assert!(offset <= d.len());
+        let n = d.len() - offset;
         let log = ceil_pow2(n as u32) as usize;
         let size = 1 << log;
-        let mut d = vec![M::identity(); 2 * size];
-        d[size..(size + n)].clone_from_slice(&v);
+        match offset.cmp(&size) {
+            Ordering::Less => {
+                d.splice(0..0, repeat_with(M::identity).take(size - offset));
+            }
+            Ordering::Equal => {}
+            Ordering::Greater => {
+                d.splice(size..offset, empty());
+            }
+        };
+        d.resize_with(size * 2, M::identity);
         let mut ret = Segtree { n, size, log, d };
         for i in (1..size).rev() {
             ret.update(i);
         }
+        // `ret.d[0]` is uninitialized and has an unknown value.
+        // This is ok as it is unused (as of writing).
         ret
     }
-}
-impl<M: Monoid> Segtree<M> {
+
     pub fn set(&mut self, mut p: usize, x: M::S) {
         assert!(p < self.n);
         p += self.size;
@@ -258,7 +286,7 @@ where
 
 #[cfg(test)]
 mod tests {
-    use crate::segtree::Max;
+    use crate::segtree::{Additive, Max};
     use crate::Segtree;
     use std::ops::{Bound::*, RangeBounds};
 
@@ -286,6 +314,14 @@ mod tests {
         check_segtree(&internal, &segtree);
     }
 
+    #[test]
+    fn test_from_iter() {
+        let it = || (1..7).map(|x| x * 4 % 11);
+        let base = it().collect::<Vec<_>>();
+        let segtree: Segtree<Max<_>> = it().collect();
+        check_segtree(&base, &segtree);
+    }
+
     //noinspection DuplicatedCode
     fn check_segtree(base: &[i32], segtree: &Segtree<Max<i32>>) {
         let n = base.len();
@@ -342,6 +378,37 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_from_vec() {
+        let v = vec![1, 2, 4];
+        let ans_124 = vec![7, 3, 4, 1, 2, 4, 0];
+        let tree = Segtree::<Additive<_>>::from_vec(v, 0);
+        assert_eq!(&tree.d[1..], &ans_124[..]);
+
+        let v = vec![1, 2, 4, 8];
+        let tree = Segtree::<Additive<_>>::from_vec(v, 0);
+        assert_eq!(&tree.d[1..], &vec![15, 3, 12, 1, 2, 4, 8][..]);
+
+        let v = vec![1, 2, 4, 8, 16];
+        let tree = Segtree::<Additive<_>>::from_vec(v, 0);
+        assert_eq!(
+            &tree.d[1..],
+            &vec![31, 15, 16, 3, 12, 16, 0, 1, 2, 4, 8, 16, 0, 0, 0][..]
+        );
+
+        let v = vec![314, 159, 265, 1, 2, 4];
+        let tree = Segtree::<Additive<_>>::from_vec(v, 3);
+        assert_eq!(&tree.d[1..], &ans_124[..]);
+
+        let v = vec![314, 159, 265, 897, 1, 2, 4];
+        let tree = Segtree::<Additive<_>>::from_vec(v, 4);
+        assert_eq!(&tree.d[1..], &ans_124[..]);
+
+        let v = vec![314, 159, 265, 897, 932, 1, 2, 4];
+        let tree = Segtree::<Additive<_>>::from_vec(v, 5);
+        assert_eq!(&tree.d[1..], &ans_124[..]);
+    }
+
     fn check(base: &[i32], segtree: &Segtree<Max<i32>>, range: impl RangeBounds<usize>) {
         let expected = base
             .iter()