competitive/algorithm/
sqrt_decomposition.rs

1#![allow(clippy::type_complexity)]
2
3use super::{Magma, Monoid, RangeBoundsExt, Unital};
4use std::ops::RangeBounds;
5
6pub trait SqrtDecomposition: Sized {
7    type M: Monoid;
8    type B;
9    fn bucket(bsize: usize) -> Self::B;
10    fn update_bucket(bucket: &mut Self::B, x: &<Self::M as Magma>::T);
11    fn update_cell(
12        bucket: &mut Self::B,
13        cell: &mut <Self::M as Magma>::T,
14        x: &<Self::M as Magma>::T,
15    );
16    fn fold_bucket(bucket: &Self::B) -> <Self::M as Magma>::T;
17    fn fold_cell(bucket: &Self::B, cell: &<Self::M as Magma>::T) -> <Self::M as Magma>::T;
18    fn sqrt_decomposition(n: usize, bucket_size: usize) -> SqrtDecompositionBuckets<Self> {
19        let mut buckets = vec![];
20        for l in (0..n).step_by(bucket_size) {
21            let bsize = (l + bucket_size).min(n) - l;
22            let x = Self::bucket(bsize);
23            buckets.push((vec![Self::M::unit(); bsize], x));
24        }
25        SqrtDecompositionBuckets {
26            n,
27            bucket_size,
28            buckets,
29            _marker: std::marker::PhantomData,
30        }
31    }
32}
33
34pub struct SqrtDecompositionBuckets<S>
35where
36    S: SqrtDecomposition,
37{
38    n: usize,
39    bucket_size: usize,
40    buckets: Vec<(Vec<<S::M as Magma>::T>, S::B)>,
41    _marker: std::marker::PhantomData<fn() -> S>,
42}
43impl<S> SqrtDecompositionBuckets<S>
44where
45    S: SqrtDecomposition,
46{
47    pub fn update_cell(&mut self, i: usize, x: <S::M as Magma>::T) {
48        let (cells, bucket) = &mut self.buckets[i / self.bucket_size];
49        let j = i % self.bucket_size;
50        S::update_cell(bucket, &mut cells[j], &x);
51    }
52    pub fn update<R>(&mut self, range: R, x: <S::M as Magma>::T)
53    where
54        R: RangeBounds<usize>,
55    {
56        let range = range.to_range_bounded(0, self.n).expect("invalid range");
57        for (i, (cells, bucket)) in self.buckets.iter_mut().enumerate() {
58            let s = i * self.bucket_size;
59            let t = s + cells.len();
60            if t <= range.start || range.end <= s {
61            } else if range.start <= s && t <= range.end {
62                S::update_bucket(bucket, &x);
63            } else {
64                for cell in &mut cells[range.start.max(s) - s..range.end.min(t) - s] {
65                    S::update_cell(bucket, cell, &x);
66                }
67            }
68        }
69    }
70    pub fn get(&self, i: usize) -> <S::M as Magma>::T {
71        let (cells, bucket) = &self.buckets[i / self.bucket_size];
72        let j = i % self.bucket_size;
73        S::fold_cell(bucket, &cells[j])
74    }
75    pub fn fold<R>(&self, range: R) -> <S::M as Magma>::T
76    where
77        R: RangeBounds<usize>,
78    {
79        let range = range.to_range_bounded(0, self.n).expect("invalid range");
80        let mut res = S::M::unit();
81        for (i, (cells, bucket)) in self.buckets.iter().enumerate() {
82            let s = i * self.bucket_size;
83            let t = s + cells.len();
84            if t <= range.start || range.end <= s {
85            } else if range.start <= s && t <= range.end {
86                <S::M as Magma>::operate_assign(&mut res, &S::fold_bucket(bucket));
87            } else {
88                for cell in &cells[range.start.max(s) - s..range.end.min(t) - s] {
89                    <S::M as Magma>::operate_assign(&mut res, &S::fold_cell(bucket, cell));
90                }
91            }
92        }
93        res
94    }
95}
competitive/algorithm/sqrt_decomposition.rs

competitive/algorithm/
sqrt_decomposition.rs
