diff --git a/src/lib.rs b/src/lib.rs index 09cc26f..0887273 100644 --- a/src/lib.rs +++ b/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, diff --git a/src/segtree.rs b/src/segtree.rs index 573b9ff..e79c74e 100644 --- a/src/segtree.rs +++ b/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(Infallible, PhantomData S>); +impl Monoid for BitwiseOr +where + S: Copy + BitOr + 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(Infallible, PhantomData S>); +impl Monoid for BitwiseAnd +where + S: Copy + BitAnd + Not + 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(Infallible, PhantomData S>); +impl Monoid for BitwiseXor +where + S: Copy + BitXor + Zero, +{ + type S = S; + fn identity() -> Self::S { + S::zero() + } + fn binary_operation(a: &Self::S, b: &Self::S) -> Self::S { + *a ^ *b + } +} + impl Default for Segtree { fn default() -> Self { Segtree::new(0) @@ -84,7 +127,27 @@ impl From> for Segtree { 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 FromIterator for Segtree { + fn from_iter>(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 Segtree { self.d[p + self.size].clone() } + pub fn get_slice(&self) -> &[M::S] { + &self.d[self.size..][..self.n] + } + pub fn prod(&self, range: R) -> M::S where R: RangeBounds,