haar_lib/algo/
sqrt_decomp.rs

1//! 平方分割
2//!
3//! # Problems
4//! - <https://onlinejudge.u-aizu.ac.jp/problems/3170>
5
6use std::ops::Range;
7
8/// 平方分割
9pub struct SqrtDecomposition {
10    size: usize,
11    block_size: usize,
12    block_num: usize,
13}
14
15impl SqrtDecomposition {
16    /// 大きさ`size`の列の平方分割の準備をする。
17    pub fn new(size: usize) -> Self {
18        let block_size = size.isqrt();
19        let block_num = size.div_ceil(block_size);
20
21        Self {
22            size,
23            block_size,
24            block_num,
25        }
26    }
27
28    /// 各ブロックで初期化を施す。
29    pub fn init(&self, mut init: impl FnMut(usize, Range<usize>)) {
30        for i in 0..self.block_num {
31            let l = i * self.block_size;
32            let r = self.size.min((i + 1) * self.block_size);
33            init(i, l..r);
34        }
35    }
36
37    /// 分割したブロック数を返す。
38    pub fn block_num(&self) -> usize {
39        self.block_num
40    }
41
42    /// `range`の範囲でのクエリを行う。
43    ///
44    /// `f`には、ブロックの番号、ブロックの範囲、ブロックの部分クエリの範囲(`None`のときブロックの全体クエリ、`Some`のときブロックの部分クエリ)が与えられる。
45    pub fn query<F>(&self, range: Range<usize>, mut f: F)
46    where
47        F: FnMut(usize, Range<usize>, Option<Range<usize>>),
48    {
49        let Range { start: l, end: r } = range;
50
51        for i in 0..self.block_num {
52            let b_l = i * self.block_size;
53            let b_r = self.size.min((i + 1) * self.block_size);
54
55            if l <= b_l && b_r <= r {
56                f(i, b_l..b_r, None);
57            } else if (b_l <= l && l < b_r) || (b_l < r && r <= b_r) {
58                f(i, b_l..b_r, Some(l.max(b_l)..r.min(b_r)));
59            }
60        }
61    }
62}