diff --git a/examples/practice2_j_segment_tree.rs b/examples/practice2_j_segment_tree.rs new file mode 100644 index 0000000..ed5d043 --- /dev/null +++ b/examples/practice2_j_segment_tree.rs @@ -0,0 +1,35 @@ +use ac_library_rs::{Max, Segtree}; +use std::io::Read; + +fn main() { + let mut buf = String::new(); + std::io::stdin().read_to_string(&mut buf).unwrap(); + let mut input = buf.split_whitespace(); + + let n: usize = input.next().unwrap().parse().unwrap(); + let q: usize = input.next().unwrap().parse().unwrap(); + let mut segtree = Segtree::>::new(n + 1); + for i in 1..=n { + segtree.set(i, input.next().unwrap().parse().unwrap()); + } + for _ in 0..q { + match input.next().unwrap().parse().unwrap() { + 1 => { + let x = input.next().unwrap().parse().unwrap(); + let v = input.next().unwrap().parse().unwrap(); + segtree.set(x, v); + } + 2 => { + let l = input.next().unwrap().parse().unwrap(); + let r: usize = input.next().unwrap().parse().unwrap(); + println!("{}", segtree.prod(l, r + 1)); + } + 3 => { + let x = input.next().unwrap().parse().unwrap(); + let v = input.next().unwrap().parse().unwrap(); + println!("{}", segtree.max_right(x, |a| a < &v)) + } + _ => {} + } + } +} diff --git a/src/internal_type_traits.rs b/src/internal_type_traits.rs index 73ba2e6..88a52a9 100644 --- a/src/internal_type_traits.rs +++ b/src/internal_type_traits.rs @@ -52,37 +52,65 @@ pub trait Integral: + fmt::Debug + fmt::Binary + fmt::Octal + + Zero + + 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; +} + +pub trait BoundedAbove { fn max_value() -> Self; } macro_rules! impl_integral { ($($ty:ty),*) => { $( - impl Integral for $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 { Self::min_value() } + } + impl BoundedAbove for $ty { #[inline] fn max_value() -> Self { Self::max_value() } } + + impl Integral for $ty {} )* }; } diff --git a/src/lib.rs b/src/lib.rs index e95c70c..d29f3d3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -28,6 +28,7 @@ pub use modint::{ ModInt998244353, Modulus, RemEuclidU32, StaticModInt, }; pub use scc::SccGraph; +pub use segtree::{Additive, 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 8b13789..2265ff1 100644 --- a/src/segtree.rs +++ b/src/segtree.rs @@ -1 +1,314 @@ +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::marker::PhantomData; +use std::ops::{Add, Mul}; +// TODO Should I split monoid-related traits to another module? +pub trait Monoid { + type S: Clone; + fn identity() -> Self::S; + fn binary_operation(a: &Self::S, b: &Self::S) -> Self::S; +} + +pub struct Max(Infallible, PhantomData S>); +impl Monoid for Max +where + S: Copy + Ord + BoundedBelow, +{ + type S = S; + fn identity() -> Self::S { + S::min_value() + } + fn binary_operation(a: &Self::S, b: &Self::S) -> Self::S { + max(*a, *b) + } +} + +pub struct Min(Infallible, PhantomData S>); +impl Monoid for Min +where + S: Copy + Ord + BoundedAbove, +{ + type S = S; + fn identity() -> Self::S { + S::max_value() + } + fn binary_operation(a: &Self::S, b: &Self::S) -> Self::S { + min(*a, *b) + } +} + +pub struct Additive(Infallible, PhantomData S>); +impl Monoid for Additive +where + S: Copy + Add + 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 Multiplicative(Infallible, PhantomData S>); +impl Monoid for Multiplicative +where + S: Copy + Mul + One, +{ + type S = S; + fn identity() -> Self::S { + S::one() + } + fn binary_operation(a: &Self::S, b: &Self::S) -> Self::S { + *a * *b + } +} + +impl Default for Segtree { + fn default() -> Self { + Segtree::new(0) + } +} +impl Segtree { + pub fn new(n: usize) -> Segtree { + vec![M::identity(); n].into() + } +} +impl From> for Segtree { + fn from(v: Vec) -> Self { + let n = v.len(); + 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); + let mut ret = Segtree { n, size, log, d }; + for i in (1..size).rev() { + ret.update(i); + } + ret + } +} +impl Segtree { + pub fn set(&mut self, mut p: usize, x: M::S) { + assert!(p < self.n); + p += self.size; + self.d[p] = x; + for i in 1..=self.log { + self.update(p >> i); + } + } + + pub fn get(&self, p: usize) -> M::S { + assert!(p < self.n); + self.d[p + self.size].clone() + } + + pub fn prod(&self, mut l: usize, mut r: usize) -> M::S { + assert!(l <= r && r <= self.n); + let mut sml = M::identity(); + let mut smr = M::identity(); + l += self.size; + r += self.size; + + while l < r { + if l & 1 != 0 { + sml = M::binary_operation(&sml, &self.d[l]); + l += 1; + } + if r & 1 != 0 { + r -= 1; + smr = M::binary_operation(&self.d[r], &smr); + } + l >>= 1; + r >>= 1; + } + + M::binary_operation(&sml, &smr) + } + + pub fn all_prod(&self) -> M::S { + self.d[1].clone() + } + + pub fn max_right(&self, mut l: usize, f: F) -> usize + where + F: Fn(&M::S) -> bool, + { + assert!(l <= self.n); + assert!(f(&M::identity())); + if l == self.n { + return self.n; + } + l += self.size; + let mut sm = M::identity(); + while { + // do + while l % 2 == 0 { + l >>= 1; + } + if !f(&M::binary_operation(&sm, &self.d[l])) { + while l < self.size { + l *= 2; + let res = M::binary_operation(&sm, &self.d[l]); + if f(&res) { + sm = res; + l += 1; + } + } + return l - self.size; + } + sm = M::binary_operation(&sm, &self.d[l]); + l += 1; + // while + { + let l = l as isize; + (l & -l) != l + } + } {} + self.n + } + + pub fn min_left(&self, mut r: usize, f: F) -> usize + where + F: Fn(&M::S) -> bool, + { + assert!(r <= self.n); + assert!(f(&M::identity())); + if r == 0 { + return 0; + } + r += self.size; + let mut sm = M::identity(); + while { + // do + r -= 1; + while r > 1 && r % 2 == 1 { + r >>= 1; + } + if !f(&M::binary_operation(&self.d[r], &sm)) { + while r < self.size { + r = 2 * r + 1; + let res = M::binary_operation(&self.d[r], &sm); + if f(&res) { + sm = res; + r -= 1; + } + } + return r + 1 - self.size; + } + sm = M::binary_operation(&self.d[r], &sm); + // while + { + let r = r as isize; + (r & -r) != r + } + } {} + 0 + } + + fn update(&mut self, k: usize) { + self.d[k] = M::binary_operation(&self.d[2 * k], &self.d[2 * k + 1]); + } +} + +// Maybe we can use this someday +// ``` +// for i in 0..=self.log { +// for j in 0..1 << i { +// print!("{}\t", self.d[(1 << i) + j]); +// } +// println!(); +// } +// ``` + +pub struct Segtree +where + M: Monoid, +{ + // variable name is _n in original library + n: usize, + size: usize, + log: usize, + d: Vec, +} + +#[cfg(test)] +mod tests { + use crate::segtree::Max; + use crate::Segtree; + + #[test] + fn test_max_segtree() { + let base = vec![3, 1, 4, 1, 5, 9, 2, 6, 5, 3]; + let n = base.len(); + let segtree: Segtree> = base.clone().into(); + check_segtree(&base, &segtree); + + let mut segtree = Segtree::>::new(n); + let mut internal = vec![i32::min_value(); n]; + for i in 0..n { + segtree.set(i, base[i]); + internal[i] = base[i]; + check_segtree(&internal, &segtree); + } + + segtree.set(6, 5); + internal[6] = 5; + check_segtree(&internal, &segtree); + + segtree.set(6, 0); + internal[6] = 0; + check_segtree(&internal, &segtree); + } + + fn check_segtree(base: &[i32], segtree: &Segtree>) { + let n = base.len(); + #[allow(clippy::needless_range_loop)] + for i in 0..n { + assert_eq!(segtree.get(i), base[i]); + } + for i in 0..=n { + for j in i..=n { + assert_eq!( + segtree.prod(i, j), + base[i..j].iter().max().copied().unwrap_or(i32::min_value()) + ); + } + } + assert_eq!( + segtree.all_prod(), + base.iter().max().copied().unwrap_or(i32::min_value()) + ); + for k in 0..=10 { + let f = |&x: &i32| x < k; + for i in 0..=n { + assert_eq!( + Some(segtree.max_right(i, f)), + (i..=n) + .filter(|&j| f(&base[i..j] + .iter() + .max() + .copied() + .unwrap_or(i32::min_value()))) + .max() + ); + } + for j in 0..=n { + assert_eq!( + Some(segtree.min_left(j, f)), + (0..=j) + .filter(|&i| f(&base[i..j] + .iter() + .max() + .copied() + .unwrap_or(i32::min_value()))) + .min() + ); + } + } + } +}