Implement From<Vec<T>> and FromIterator for FenwickTree

examples/library-checker-static-range-sum.rs

@@ -15,7 +15,7 @@ fn main() {
         lrs: [(usize, usize); q],
     }
 
-    let mut fenwick = FenwickTree::new(n, 0);
+    let mut fenwick = FenwickTree::<u64>::new(n);
     for (i, a) in r#as.into_iter().enumerate() {
         fenwick.add(i, a);
     }

src/fenwicktree.rs

@@ -1,20 +1,22 @@
+use std::{iter::FromIterator, ops::AddAssign};
+
+use crate::num_traits::Zero;
+
 // Reference: https://en.wikipedia.org/wiki/Fenwick_tree
 pub struct FenwickTree<T> {
     n: usize,
     ary: Vec<T>,
-    e: T,
 }
 
-impl<T: Clone + std::ops::AddAssign<T>> FenwickTree<T> {
-    pub fn new(n: usize, e: T) -> Self {
+impl<T: Clone + AddAssign<T> + Zero> FenwickTree<T> {
+    pub fn new(n: usize) -> Self {
         FenwickTree {
             n,
-            ary: vec![e.clone(); n],
-            e,
+            ary: vec![T::zero(); n],
         }
     }
     pub fn accum(&self, mut idx: usize) -> T {
-        let mut sum = self.e.clone();
+        let mut sum = T::zero();
         while idx > 0 {
             sum += self.ary[idx - 1].clone();
             idx &= idx - 1;
@@ -24,7 +26,7 @@ impl<T: Clone + std::ops::AddAssign<T>> FenwickTree<T> {
     /// performs data[idx] += val;
     pub fn add<U: Clone>(&mut self, mut idx: usize, val: U)
     where
-        T: std::ops::AddAssign<U>,
+        T: AddAssign<U>,
     {
         let n = self.n;
         idx += 1;
@@ -41,14 +43,31 @@ impl<T: Clone + std::ops::AddAssign<T>> FenwickTree<T> {
         self.accum(r) - self.accum(l)
     }
 }
+impl<T: Clone + AddAssign<T>> From<Vec<T>> for FenwickTree<T> {
+    fn from(mut ary: Vec<T>) -> Self {
+        for i in 1..=ary.len() {
+            let j = i + (i & i.wrapping_neg());
+            if j <= ary.len() {
+                let add = ary[i - 1].clone();
+                ary[j - 1] += add;
+            }
+        }
+        Self { n: ary.len(), ary }
+    }
+}
+impl<T: Clone + AddAssign<T>> FromIterator<T> for FenwickTree<T> {
+    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
+        iter.into_iter().collect::<Vec<_>>().into()
+    }
+}
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
     fn fenwick_tree_works() {
-        let mut bit = FenwickTree::new(5, 0i64);
+        let mut bit = FenwickTree::<i64>::new(5);
         // [1, 2, 3, 4, 5]
         for i in 0..5 {
             bit.add(i, i as i64 + 1);
@@ -57,4 +76,15 @@ mod tests {
         assert_eq!(bit.sum(0, 4), 10);
         assert_eq!(bit.sum(1, 3), 5);
     }
+
+    #[test]
+    fn from_iter_works() {
+        let tree = FenwickTree::from_iter(vec![1, 2, 3, 4, 5].iter().map(|x| x * 2));
+        let internal = vec![2, 4, 6, 8, 10];
+        for j in 0..=internal.len() {
+            for i in 0..=j {
+                assert_eq!(tree.sum(i, j), internal[i..j].iter().sum::<i32>());
+            }
+        }
+    }
 }

src/internal_type_traits.rs

@@ -52,25 +52,13 @@ pub trait Integral:
     + fmt::Debug
     + fmt::Binary
     + fmt::Octal
-    + Zero
-    + One
+    + crate::num_traits::Zero
+    + crate::num_traits::One
     + BoundedBelow
     + BoundedAbove
 {
 }
 
-/// Class that has additive identity element
-pub trait Zero {
-    /// The additive identity element
-    fn zero() -> Self;
-}
-
-/// Class that has multiplicative identity element
-pub trait One {
-    /// The multiplicative identity element
-    fn one() -> Self;
-}
-
 pub trait BoundedBelow {
     fn min_value() -> Self;
 }
@@ -82,20 +70,6 @@ pub trait BoundedAbove {
 macro_rules! impl_integral {
     ($($ty:ty),*) => {
         $(
-            impl Zero for $ty {
-                #[inline]
-                fn zero() -> Self {
-                    0
-                }
-            }
-
-            impl One for $ty {
-                #[inline]
-                fn one() -> Self {
-                    1
-                }
-            }
-
             impl BoundedBelow for $ty {
                 #[inline]
                 fn min_value() -> Self {

src/lib.rs

@@ -11,6 +11,8 @@ pub mod segtree;
 pub mod string;
 pub mod twosat;
 
+pub mod num_traits;
+
 pub(crate) mod internal_bit;
 pub(crate) mod internal_math;
 pub(crate) mod internal_queue;

src/num_traits.rs

@@ -0,0 +1,43 @@
+/// A type that has an additive identity element.
+pub trait Zero {
+    /// The additive identity element.
+    fn zero() -> Self;
+}
+
+/// A type that has a multiplicative identity element.
+pub trait One {
+    /// The multiplicative identity element.
+    fn one() -> Self;
+}
+
+macro_rules! impl_zero {
+    ($zero: literal, $one: literal: $($t: ty),*) => {
+        $(
+            impl Zero for $t {
+                fn zero() -> Self {
+                    $zero
+                }
+            }
+
+            impl One for $t {
+                fn one() -> Self {
+                    $one
+                }
+            }
+        )*
+    };
+}
+impl_zero!(0, 1: usize, u8, u16, u32, u64, u128);
+impl_zero!(0, 1: isize, i8, i16, i32, i64, i128);
+impl_zero!(0.0, 1.0: f32, f64);
+
+impl<T: Zero> Zero for core::num::Wrapping<T> {
+    fn zero() -> Self {
+        Self(T::zero())
+    }
+}
+impl<T: One> One for core::num::Wrapping<T> {
+    fn one() -> Self {
+        Self(T::one())
+    }
+}

src/segtree.rs

@@ -1,5 +1,6 @@
 use crate::internal_bit::ceil_pow2;
-use crate::internal_type_traits::{BoundedAbove, BoundedBelow, One, Zero};
+use crate::internal_type_traits::{BoundedAbove, BoundedBelow};
+use crate::num_traits::{One, Zero};
 use std::cmp::{max, min};
 use std::convert::Infallible;
 use std::marker::PhantomData;