Merge pull request #115 from mizar/segtree

Segtree: Monoid BitOrOper/BitAndOper/BitXorOper, FromIterator, get_slice

src/lib.rs

@@ -29,7 +29,9 @@ pub use modint::{
     ModInt1000000007, ModInt998244353, Modulus, RemEuclidU32, StaticModInt,
 };
 pub use scc::SccGraph;
-pub use segtree::{Additive, Max, Min, Monoid, Multiplicative, Segtree};
+pub use segtree::{
+    Additive, BitwiseAnd, BitwiseOr, BitwiseXor, Max, Min, Monoid, Multiplicative, Segtree,
+};
 pub use string::{
     lcp_array, lcp_array_arbitrary, suffix_array, suffix_array_arbitrary, suffix_array_manual,
     z_algorithm, z_algorithm_arbitrary,

src/segtree.rs

@@ -2,8 +2,9 @@ use crate::internal_bit::ceil_pow2;
 use crate::internal_type_traits::{BoundedAbove, BoundedBelow, One, Zero};
 use std::cmp::{max, min};
 use std::convert::Infallible;
+use std::iter::FromIterator;
 use std::marker::PhantomData;
-use std::ops::{Add, Bound, Mul, RangeBounds};
+use std::ops::{Add, BitAnd, BitOr, BitXor, Bound, Mul, Not, RangeBounds};
 
 // TODO Should I split monoid-related traits to another module?
 pub trait Monoid {
@@ -68,6 +69,48 @@ where
     }
 }
 
+pub struct BitwiseOr<S>(Infallible, PhantomData<fn() -> S>);
+impl<S> Monoid for BitwiseOr<S>
+where
+    S: Copy + BitOr<Output = S> + Zero,
+{
+    type S = S;
+    fn identity() -> Self::S {
+        S::zero()
+    }
+    fn binary_operation(a: &Self::S, b: &Self::S) -> Self::S {
+        *a | *b
+    }
+}
+
+pub struct BitwiseAnd<S>(Infallible, PhantomData<fn() -> S>);
+impl<S> Monoid for BitwiseAnd<S>
+where
+    S: Copy + BitAnd<Output = S> + Not<Output = S> + Zero,
+{
+    type S = S;
+    fn identity() -> Self::S {
+        !S::zero()
+    }
+    fn binary_operation(a: &Self::S, b: &Self::S) -> Self::S {
+        *a & *b
+    }
+}
+
+pub struct BitwiseXor<S>(Infallible, PhantomData<fn() -> S>);
+impl<S> Monoid for BitwiseXor<S>
+where
+    S: Copy + BitXor<Output = S> + Zero,
+{
+    type S = S;
+    fn identity() -> Self::S {
+        S::zero()
+    }
+    fn binary_operation(a: &Self::S, b: &Self::S) -> Self::S {
+        *a ^ *b
+    }
+}
+
 impl<M: Monoid> Default for Segtree<M> {
     fn default() -> Self {
         Segtree::new(0)
@@ -84,7 +127,27 @@ impl<M: Monoid> From<Vec<M::S>> for Segtree<M> {
         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);
+        d[size..][..n].clone_from_slice(&v);
+        let mut ret = Segtree { n, size, log, d };
+        for i in (1..size).rev() {
+            ret.update(i);
+        }
+        ret
+    }
+}
+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(
+            std::iter::repeat_with(M::identity)
+                .take(size)
+                .chain(iter)
+                .chain(std::iter::repeat_with(M::identity).take(size - n)),
+        );
         let mut ret = Segtree { n, size, log, d };
         for i in (1..size).rev() {
             ret.update(i);
@@ -107,6 +170,10 @@ impl<M: Monoid> Segtree<M> {
         self.d[p + self.size].clone()
     }
 
+    pub fn get_slice(&self) -> &[M::S] {
+        &self.d[self.size..][..self.n]
+    }
+
     pub fn prod<R>(&self, range: R) -> M::S
     where
         R: RangeBounds<usize>,