haar_lib/ds/
range_search_tree.rs

1//! 領域内の点を列挙する
2//!
3//! # Problems
4//! - [AOJ DSL 2_C: Range Search(kD Tree)](https://onlinejudge.u-aizu.ac.jp/courses/library/3/DSL/2/DSL_2_C)
5
6use crate::algo::{bsearch_slice::BinarySearch, merge::merge};
7
8/// Range search tree
9pub struct RangeSearchTree<Index> {
10    size: usize,
11    cxs: Vec<Index>,
12    data: Vec<Vec<(Index, usize)>>,
13}
14
15impl<Index> RangeSearchTree<Index>
16where
17    Index: Copy + Ord,
18{
19    /// $[sx, tx), [sy, ty)$の矩形内部にある点を列挙する。
20    pub fn search(
21        &self,
22        (sx, sy): (Index, Index),
23        (tx, ty): (Index, Index),
24    ) -> Vec<(Index, Index)> {
25        assert!(sx < tx);
26        assert!(sy < ty);
27
28        let mut ret = vec![];
29        let mut l = self.cxs.lower_bound(&sx) + self.size / 2;
30        let mut r = self.cxs.lower_bound(&tx) + self.size / 2;
31
32        while l < r {
33            if (r & 1) != 0 {
34                r -= 1;
35                let a = &self.data[r];
36
37                let i = a.lower_bound(&(sy, 0));
38
39                for &(y, x) in a.iter().skip(i).take_while(|(y, _)| *y < ty) {
40                    ret.push((self.cxs[x], y));
41                }
42            }
43
44            if (l & 1) != 0 {
45                let a = &self.data[l];
46                l += 1;
47
48                let i = a.lower_bound(&(sy, 0));
49
50                for &(y, x) in a.iter().skip(i).take_while(|(y, _)| *y < ty) {
51                    ret.push((self.cxs[x], y));
52                }
53            }
54
55            l >>= 1;
56            r >>= 1;
57        }
58
59        ret
60    }
61}
62
63/// [`RangeSearchTree`]を構築するための構造体。
64#[derive(Clone, Default)]
65pub struct RangeSearchTreeBuilder<Index> {
66    size: usize,
67    xs: Vec<Index>,
68    ys: Vec<Index>,
69}
70
71impl<Index> RangeSearchTreeBuilder<Index>
72where
73    Index: Copy + Ord,
74{
75    /// 空の[`RangeSearchTreeBuilder`]を用意する。
76    pub fn new() -> Self {
77        Self {
78            size: 0,
79            xs: vec![],
80            ys: vec![],
81        }
82    }
83
84    /// 点`(x, y)`を登録する。
85    pub fn add(&mut self, x: Index, y: Index) {
86        self.size += 1;
87        self.xs.push(x);
88        self.ys.push(y);
89    }
90
91    /// [`RangeSearchTree`]を構築する。
92    pub fn build(self) -> RangeSearchTree<Index> {
93        let mut cxs = self.xs.clone();
94        cxs.sort_unstable();
95        cxs.dedup();
96
97        let m = cxs.len();
98        let size = m.next_power_of_two() * 2;
99
100        let mut data: Vec<Vec<(Index, usize)>> = vec![vec![]; size];
101
102        for i in 0..self.size {
103            let j = cxs.lower_bound(&self.xs[i]);
104            data[size / 2 + j].push((self.ys[i], j));
105        }
106
107        for item in data.iter_mut().take(size).skip(size / 2) {
108            item.sort_unstable();
109        }
110
111        for i in (1..size / 2).rev() {
112            data[i] = merge(data[i << 1].clone(), data[(i << 1) | 1].clone());
113        }
114
115        RangeSearchTree { size, cxs, data }
116    }
117}
haar_lib/ds/range_search_tree.rs

haar_lib/ds/
range_search_tree.rs
