competitive/data_structure/
kdtree.rs

1use std::ops::Range;
2
3pub struct Static2DTree<T, U, V>
4where
5    T: Ord,
6    U: Ord,
7{
8    data: Vec<(T, U, V)>,
9}
10impl<T, U, V> Static2DTree<T, U, V>
11where
12    T: Ord,
13    U: Ord,
14{
15    pub fn new<I>(data: I) -> Self
16    where
17        I: IntoIterator<Item = (T, U, V)>,
18    {
19        let mut data: Vec<_> = data.into_iter().collect();
20        let n = data.len();
21        Self::build(&mut data, 0, n, 0);
22        Self { data }
23    }
24    fn build(data: &mut [(T, U, V)], l: usize, r: usize, depth: usize) {
25        if l < r {
26            let m = l + (r - l) / 2;
27            if depth % 2 == 0 {
28                data[l..r].sort_by(|p, q| p.0.cmp(&q.0));
29            } else {
30                data[l..r].sort_by(|p, q| p.1.cmp(&q.1));
31            }
32            Self::build(data, l, m, depth + 1);
33            Self::build(data, m + 1, r, depth + 1);
34        }
35    }
36    pub fn range(&self, range1: Range<T>, range2: Range<U>) -> Vec<&V> {
37        let mut res = vec![];
38        self.range_inner(&range1, &range2, 0, self.data.len(), 0, &mut res);
39        res
40    }
41    fn range_inner<'a>(
42        &'a self,
43        range1: &Range<T>,
44        range2: &Range<U>,
45        l: usize,
46        r: usize,
47        depth: usize,
48        res: &mut Vec<&'a V>,
49    ) {
50        if l < r {
51            let m = l + (r - l) / 2;
52            let (t, u, v) = &self.data[m];
53            if range1.contains(t) && range2.contains(u) {
54                res.push(v);
55            }
56            if if depth % 2 == 0 {
57                &range1.start <= t
58            } else {
59                &range2.start <= u
60            } {
61                self.range_inner(range1, range2, l, m, depth + 1, res);
62            }
63            if if depth % 2 == 0 {
64                t < &range1.end
65            } else {
66                u < &range2.end
67            } {
68                self.range_inner(range1, range2, m + 1, r, depth + 1, res);
69            }
70        }
71    }
72}
competitive/data_structure/kdtree.rs

competitive/data_structure/
kdtree.rs
