2.5.4 K-d Tree for 2D Range Query

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

/*

Maintain a static set of two-dimensional points while supporting rectangle reporting queries. A k-d
tree recursively splits points by alternating coordinates and stores a bounding box for each
subtree, allowing whole subtrees to be accepted or pruned during a query.

This implementation uses `std::pair` to represent points, requiring `operator<` to be defined on the
numeric template type.

Use this for static point-reporting queries when O(n) space and good average performance are more
important than a strict worst-case guarantee. Use the 2D range tree instead when adversarial point
sets or query rectangles are expected and the extra O(n log n) space is acceptable.

- `RangeKDTree<T>(lo, hi)` constructs a set of `std::pair` points from the half-open
  forward-iterator range $[`lo`, `hi`)$.
- `query(x1, y1, x2, y2, f)` calls the function `f(p)` on each point whose $x$-coordinate is in
  $[`x1`, `x2`]$ and whose $y$-coordinate is in $[`y1`, `y2`]$. The argument `p` is the point as an
  `std::pair`.

Time Complexity:
- O(n log n) per call to the constructor, where $n$ is the number of points.
- O(sqrt(n) + m) per call to `query()`, where $m$ is the number of points that are reported by the
  query.

Space Complexity:
- O(n) for storage of the points.
- O(log n) auxiliary stack space for `query()`.

*/

#include <algorithm>
#include <cassert>
#include <utility>
#include <vector>

template<typename T>
class RangeKDTree {
  std::vector<std::pair<T, T>> tree, minp, maxp;
  std::vector<int> l_index, h_index;

  void build(int lo, int hi, bool div_x) {
    if (lo >= hi) {
      return;
    }
    int mid = lo + (hi - lo) / 2;
    std::nth_element(
        tree.begin() + lo, tree.begin() + mid, tree.begin() + hi,
        [div_x](const auto &a, const auto &b) {
          return div_x ? a.first < b.first : a.second < b.second;
        }
    );
    l_index[mid] = lo;
    h_index[mid] = hi;
    minp[mid].first = maxp[mid].first = tree[lo].first;
    minp[mid].second = maxp[mid].second = tree[lo].second;
    for (int i = lo + 1; i < hi; i++) {
      minp[mid].first = std::min(minp[mid].first, tree[i].first);
      minp[mid].second = std::min(minp[mid].second, tree[i].second);
      maxp[mid].first = std::max(maxp[mid].first, tree[i].first);
      maxp[mid].second = std::max(maxp[mid].second, tree[i].second);
    }
    build(lo, mid, !div_x);
    build(mid + 1, hi, !div_x);
  }

  template<typename Fn>
  void query(int lo, int hi, const T &x1, const T &y1, const T &x2, const T &y2, Fn &f) {
    if (lo >= hi) {
      return;
    }
    int mid = lo + (hi - lo) / 2;
    T ax = minp[mid].first, ay = minp[mid].second;
    T bx = maxp[mid].first, by = maxp[mid].second;
    if (x2 < ax || bx < x1 || y2 < ay || by < y1) {
      return;
    }
    if (!(ax < x1 || x2 < bx || ay < y1 || y2 < by)) {
      for (int i = l_index[mid]; i < h_index[mid]; i++) {
        f(tree[i]);
      }
      return;
    }
    query(lo, mid, x1, y1, x2, y2, f);
    query(mid + 1, hi, x1, y1, x2, y2, f);
    if (tree[mid].first < x1 || x2 < tree[mid].first || tree[mid].second < y1 ||
        y2 < tree[mid].second) {
      return;
    }
    f(tree[mid]);
  }

 public:
  template<typename It>
  RangeKDTree(It lo, It hi)
      : tree(lo, hi),
        minp(tree.size()),
        maxp(tree.size()),
        l_index(tree.size()),
        h_index(tree.size()) {
    build(0, static_cast<int>(tree.size()), true);
  }

  template<typename Fn>
  void query(const T &x1, const T &y1, const T &x2, const T &y2, Fn f) {
    assert(!(x2 < x1) && !(y2 < y1));
    query(0, static_cast<int>(tree.size()), x1, y1, x2, y2, f);
  }
};

/*** Example Usage ***/

#include <algorithm>
#include <cassert>
using namespace std;

int main() {
  vector<pair<int, int>> v{{1, 4},  {5, 4},  {2, 2},   {3, 1},   {6, -5},
                           {5, -1}, {3, -3}, {-1, -2}, {-1, -1}, {2, -1}};
  RangeKDTree<int> t(v.begin(), v.end());
  vector<pair<int, int>> got;
  auto collect = [&](const pair<int, int> &p) { got.push_back(p); };
  t.query(-1, -1, 2, 5, collect);
  sort(got.begin(), got.end());
  assert((got == vector<pair<int, int>>{{-1, -1}, {1, 4}, {2, -1}, {2, 2}}));
  got.clear();
  t.query(1, 1, 4, 8, collect);
  sort(got.begin(), got.end());
  assert((got == vector<pair<int, int>>{{1, 4}, {2, 2}, {3, 1}}));
  return 0;
}