sqrt_bucket/
lib.rs

1use std::ops::{Range, RangeBounds};
2
3use usize_bounds::UsizeBounds;
4
5pub struct SqrtBucket<T, S, Ff, Fr, Fq> {
6    buf: Vec<T>,
7    summary: Vec<S>,
8    force: Ff,
9    reduce: Fr,
10    query: Fq,
11    bucket_size: usize,
12}
13
14pub enum BucketBorrow<'a, T, S> {
15    Slice(&'a mut [T]),
16    Summary(&'a mut S),
17}
18
19impl<T, S, Ff, Fr, Fq> SqrtBucket<T, S, Ff, Fr, Fq>
20where
21    Fr: FnMut(&[T]) -> S,
22    Ff: FnMut(&S, &mut [T]),
23{
24    pub fn new(buf: Vec<T>, force: Ff, reduce: Fr, query: Fq) -> Self {
25        Self::new_with_bucket_size(buf, force, reduce, query, 384)
26    }
27
28    pub fn new_with_bucket_size(
29        buf: Vec<T>,
30        force: Ff,
31        mut reduce: Fr,
32        query: Fq,
33        bucket_size: usize,
34    ) -> Self {
35        let summary: Vec<S> =
36            buf.chunks(bucket_size).map(&mut reduce).collect();
37        Self { buf, summary, force, reduce, query, bucket_size }
38    }
39
40    pub fn query<Q, B, R>(&mut self, range: B, args: Q) -> R
41    where
42        B: RangeBounds<usize>,
43        Fq: for<'a> FnMut(&mut [BucketBorrow<'a, T, S>], Q) -> R,
44    {
45        let n = self.buf.len();
46        let b = self.bucket_size;
47        let Range { start, end } = range.to_range(n);
48
49        let mut borrowed = vec![];
50        let mut affected = vec![];
51        for ((l, chunk_i), summary_i) in
52            (0..n).step_by(b).zip(self.buf.chunks_mut(b)).zip(&mut self.summary)
53        {
54            let i = l / b;
55            let r = n.min(l + b);
56            if r <= start {
57                continue;
58            }
59            if end <= l {
60                break;
61            }
62            if start <= l && r <= end {
63                borrowed.push(BucketBorrow::Summary(summary_i));
64            } else {
65                (self.force)(summary_i, chunk_i);
66                affected.push((i, (l, r)));
67                let jl = l.max(start) - l;
68                let jr = r.min(end) - l;
69                borrowed.push(BucketBorrow::Slice(&mut chunk_i[jl..jr]));
70            }
71        }
72        let res = (self.query)(&mut borrowed, args);
73        for (i, (l, r)) in affected {
74            self.summary[i] = (self.reduce)(&mut self.buf[l..r]);
75        }
76        res
77    }
78}
79
80// TODO: doc (cf. <https://atcoder.jp/contests/abc322/submissions/53952026>)
sqrt_bucket/lib.rs

sqrt_bucket/
lib.rs
