From a47e108cc72cd69d24893b024ff060d2b80f2ade Mon Sep 17 00:00:00 2001
From: rsk0315 <rsk0315@icloud.com>
Date: Mon, 27 May 2024 21:20:36 +0900
Subject: [PATCH] a

---
 nekolib-src/ds/Cargo.toml             |  1 +
 nekolib-src/ds/sqrt_bucket/Cargo.toml |  7 +++
 nekolib-src/ds/sqrt_bucket/src/lib.rs | 80 +++++++++++++++++++++++++++
 nekolib-src/ds/src/lib.rs             |  1 +
 4 files changed, 89 insertions(+)
 create mode 100644 nekolib-src/ds/sqrt_bucket/Cargo.toml
 create mode 100644 nekolib-src/ds/sqrt_bucket/src/lib.rs

diff --git a/nekolib-src/ds/Cargo.toml b/nekolib-src/ds/Cargo.toml
index 64de6544b3..b2b54bf2be 100644
--- a/nekolib-src/ds/Cargo.toml
+++ b/nekolib-src/ds/Cargo.toml
@@ -15,3 +15,4 @@ rs01_dict = { path = "rs01_dict" }
 wavelet_matrix = { path = "wavelet_matrix" }
 fibonacci_heap = { path = "fibonacci_heap" }
 btree_seq = { path = "btree_seq" }
+sqrt_bucket = { path = "sqrt_bucket" }
diff --git a/nekolib-src/ds/sqrt_bucket/Cargo.toml b/nekolib-src/ds/sqrt_bucket/Cargo.toml
new file mode 100644
index 0000000000..8bc706bd98
--- /dev/null
+++ b/nekolib-src/ds/sqrt_bucket/Cargo.toml
@@ -0,0 +1,7 @@
+[package]
+name = "sqrt_bucket"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+usize_bounds = { path = "../../ops/usize_bounds" }
diff --git a/nekolib-src/ds/sqrt_bucket/src/lib.rs b/nekolib-src/ds/sqrt_bucket/src/lib.rs
new file mode 100644
index 0000000000..6a64f58069
--- /dev/null
+++ b/nekolib-src/ds/sqrt_bucket/src/lib.rs
@@ -0,0 +1,80 @@
+use std::ops::{Range, RangeBounds};
+
+use usize_bounds::UsizeBounds;
+
+pub struct SqrtBucket<T, S, Ff, Fr, Fq> {
+    buf: Vec<T>,
+    summary: Vec<S>,
+    force: Ff,
+    reduce: Fr,
+    query: Fq,
+    bucket_size: usize,
+}
+
+pub enum BucketBorrow<'a, T, S> {
+    Slice(&'a mut [T]),
+    Summary(&'a mut S),
+}
+
+impl<T, S, Ff, Fr, Fq> SqrtBucket<T, S, Ff, Fr, Fq>
+where
+    Fr: FnMut(&[T]) -> S,
+    Ff: FnMut(&S, &mut [T]),
+{
+    pub fn new(buf: Vec<T>, force: Ff, reduce: Fr, query: Fq) -> Self {
+        Self::new_with_bucket_size(buf, force, reduce, query, 384)
+    }
+
+    pub fn new_with_bucket_size(
+        buf: Vec<T>,
+        force: Ff,
+        mut reduce: Fr,
+        query: Fq,
+        bucket_size: usize,
+    ) -> Self {
+        let summary: Vec<S> =
+            buf.chunks(bucket_size).map(&mut reduce).collect();
+        Self { buf, summary, force, reduce, query, bucket_size }
+    }
+
+    pub fn query<Q, B, R>(&mut self, range: B, args: Q) -> R
+    where
+        B: RangeBounds<usize>,
+        Fq: for<'a> FnMut(&mut [BucketBorrow<'a, T, S>], Q) -> R,
+    {
+        let n = self.buf.len();
+        let b = self.bucket_size;
+        let Range { start, end } = range.to_range(n);
+
+        let mut borrowed = vec![];
+        let mut affected = vec![];
+        for ((l, chunk_i), summary_i) in
+            (0..n).step_by(b).zip(self.buf.chunks_mut(b)).zip(&mut self.summary)
+        {
+            let i = l / b;
+            let r = n.min(l + b);
+            if r <= start {
+                continue;
+            }
+            if end <= l {
+                break;
+            }
+            if start <= l && r <= end {
+                borrowed.push(BucketBorrow::Summary(summary_i));
+            } else {
+                (self.force)(summary_i, chunk_i);
+                affected.push((i, (l, r)));
+                let jl = l.max(start) - l;
+                let jr = r.min(end) - l;
+                borrowed.push(BucketBorrow::Slice(&mut chunk_i[jl..jr]));
+            }
+        }
+        let res = (self.query)(&mut borrowed, args);
+        for (i, (l, r)) in affected {
+            self.summary[i] = (self.reduce)(&mut self.buf[l..r]);
+        }
+        res
+    }
+}
+
+// TODO: doc (cf. <https://atcoder.jp/contests/abc322/submissions/53952026>)
diff --git a/nekolib-src/ds/src/lib.rs b/nekolib-src/ds/src/lib.rs
index e54b237015..48d7ec4a8d 100644
--- a/nekolib-src/ds/src/lib.rs
+++ b/nekolib-src/ds/src/lib.rs
@@ -6,6 +6,7 @@ doc_inline_reexport! {
     foldable_deque,
     foldable_queue,
     rs01_dict,
+    sqrt_bucket,
     union_find,
     vec_segtree,
     wavelet_matrix,