diff --git a/examples/library-checker-static-range-sum.rs b/examples/library-checker-static-range-sum.rs index 4c5cc03..54ee5bc 100644 --- a/examples/library-checker-static-range-sum.rs +++ b/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::::new(n); for (i, a) in r#as.into_iter().enumerate() { fenwick.add(i, a); } diff --git a/expand.py b/expand.py index 283d6fb..8f04fc2 100755 --- a/expand.py +++ b/expand.py @@ -32,11 +32,12 @@ output_header = '//https://github.com/rust-lang-ja/ac-library-rs\n' opt_list = ['help', 'all', 'output='] output_list_all = ('convolution', 'dsu', 'fenwicktree', 'lazysegtree', 'math', - 'maxflow', 'mincostflow', 'modint', 'scc', 'segtree', - 'string', 'twosat', + 'maxflow', 'mincostflow', 'modint', 'num_traits', 'scc', + 'segtree', 'string', 'twosat', 'internal_bit', 'internal_math', 'internal_queue', 'internal_scc', 'internal_type_traits',) dependency_list = {'convolution': ('internal_bit', 'modint',), + 'fenwicktree': ('num_traits',), 'lazysegtree': ('internal_bit', 'segtree'), 'math': ('internal_math',), 'maxflow': ('internal_type_traits', 'internal_queue',), @@ -44,7 +45,8 @@ 'modint': ('internal_math',), 'scc': ('internal_scc',), 'segtree': ('internal_bit', 'internal_type_traits',), - 'twosat': ('internal_scc',), } + 'twosat': ('internal_scc',), + 'internal_type_traits': ('num_traits',), } src_path = 'src/' output_path = None diff --git a/src/fenwicktree.rs b/src/fenwicktree.rs index cfcf32a..3ecc1af 100644 --- a/src/fenwicktree.rs +++ b/src/fenwicktree.rs @@ -1,22 +1,25 @@ -use std::ops::{Bound, RangeBounds}; +use std::{ + iter::FromIterator, + ops::{AddAssign, Bound, RangeBounds}, +}; + +use crate::num_traits::Zero; // Reference: https://en.wikipedia.org/wiki/Fenwick_tree pub struct FenwickTree { n: usize, ary: Vec, - e: T, } -impl> FenwickTree { - pub fn new(n: usize, e: T) -> Self { +impl + Zero> FenwickTree { + 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; @@ -26,7 +29,7 @@ impl> FenwickTree { /// performs data[idx] += val; pub fn add(&mut self, mut idx: usize, val: U) where - T: std::ops::AddAssign, + T: AddAssign, { let n = self.n; idx += 1; @@ -54,6 +57,23 @@ impl> FenwickTree { self.accum(r) - self.accum(l) } } +impl> From> for FenwickTree { + fn from(mut ary: Vec) -> 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> FromIterator for FenwickTree { + fn from_iter>(iter: I) -> Self { + iter.into_iter().collect::>().into() + } +} #[cfg(test)] mod tests { @@ -62,7 +82,7 @@ mod tests { #[test] fn fenwick_tree_works() { - let mut bit = FenwickTree::new(5, 0i64); + let mut bit = FenwickTree::::new(5); // [1, 2, 3, 4, 5] for i in 0..5 { bit.add(i, i as i64 + 1); @@ -78,4 +98,15 @@ mod tests { assert_eq!(bit.sum(1..=3), 9); assert_eq!(bit.sum((Excluded(0), Included(2))), 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::()); + } + } + } } diff --git a/src/internal_type_traits.rs b/src/internal_type_traits.rs index 88a52a9..9b00aec 100644 --- a/src/internal_type_traits.rs +++ b/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 { diff --git a/src/lib.rs b/src/lib.rs index 09cc26f..1a8deb4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -11,6 +11,8 @@ pub mod segtree; pub mod string; pub mod twosat; +pub mod num_traits; + mod internal_bit; mod internal_math; mod internal_queue; diff --git a/src/num_traits.rs b/src/num_traits.rs new file mode 100644 index 0000000..5fa6a23 --- /dev/null +++ b/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 Zero for core::num::Wrapping { + fn zero() -> Self { + Self(T::zero()) + } +} +impl One for core::num::Wrapping { + fn one() -> Self { + Self(T::one()) + } +} diff --git a/src/segtree.rs b/src/segtree.rs index 573b9ff..90b0037 100644 --- a/src/segtree.rs +++ b/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;