2.6.3 2D Range Tree

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
131
132
133
134
135

/*

Maintain a static collection of two-dimensional points while supporting rectangle reporting queries.
A 2D range tree sorts the points by $x$ and stores at each tree node its subrange of points sorted
by $y$, merged from its children's lists as in a merge sort tree. A query decomposes its $x$-range
into O(log n) canonical nodes and binary searches their $y$-sorted lists.

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 guaranteed worst-case bounds are more important
than memory. Compared with the k-d tree in the next section, it uses O(n log n) rather than O(n)
space but improves the worst-case query time from O(sqrt(n) + m) to O(log^2 n + m).

- `RangeTree<T>(lo, hi)` constructs a collection of `std::pair` points from the half-open
  forward-iterator range $[`lo`, `hi`)$.
- `query(x1, y1, x2, y2)` returns (`i`, `x`, `y`) tuples for all points in the closed rectangle
  $[`x1`, `x2`] \times [`y1`, `y2`]$, where `i` is the point's 0-based index in the original range.

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

Space Complexity:
- O(n log n) for storage of the points.
- O(log n) auxiliary stack space for `query()`.
- O(m) for the vector returned by `query()`, where $m$ is the number of reported points.

*/

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

template<typename T>
class RangeTree {
  struct IndexedPoint {
    T x, y;
    int original_index;
  };

  std::vector<IndexedPoint> points;
  std::vector<std::vector<std::pair<int, T>>> cols;

  void build(int n, int lo, int hi) {
    if (!(points[lo].x < points[hi].x)) {
      for (int i = lo; i <= hi; i++) {
        cols[n].emplace_back(i, points[i].y);
      }
      return;
    }
    int l = n * 2 + 1, r = n * 2 + 2, mid = lo + (hi - lo) / 2;
    build(l, lo, mid);
    build(r, mid + 1, hi);
    cols[n].resize(cols[l].size() + cols[r].size());
    std::merge(
        cols[l].begin(), cols[l].end(), cols[r].begin(), cols[r].end(), cols[n].begin(),
        [](const auto &a, const auto &b) { return a.second < b.second; }
    );
  }

  void query(
      int n, int lo, int hi, const T &x1, const T &y1, const T &x2, const T &y2,
      std::vector<std::tuple<int, T, T>> &res
  ) const {
    if (points[hi].x < x1 || x2 < points[lo].x) {
      return;
    }
    if (!(points[lo].x < x1 || x2 < points[hi].x)) {
      if (!cols[n].empty() && !(y2 < y1)) {
        auto it =
            std::lower_bound(cols[n].begin(), cols[n].end(), y1, [](const auto &a, const T &value) {
              return a.second < value;
            });
        for (; it != cols[n].end() && !(y2 < it->second); ++it) {
          const IndexedPoint &p = points[it->first];
          res.emplace_back(p.original_index, p.x, p.y);
        }
      }
    } else if (lo != hi) {
      int mid = lo + (hi - lo) / 2;
      query(n * 2 + 1, lo, mid, x1, y1, x2, y2, res);
      query(n * 2 + 2, mid + 1, hi, x1, y1, x2, y2, res);
    }
  }

 public:
  template<typename It>
  RangeTree(It lo, It hi) {
    int n = std::distance(lo, hi);
    assert(n > 0);
    points.reserve(n);
    int index = 0;
    for (It it = lo; it != hi; ++it) {
      points.push_back({it->first, it->second, index++});
    }
    cols.resize(4 * n + 1);
    std::sort(points.begin(), points.end(), [](const IndexedPoint &a, const IndexedPoint &b) {
      return a.x < b.x || (!(b.x < a.x) && a.y < b.y);
    });
    build(0, 0, n - 1);
  }

  std::vector<std::tuple<int, T, T>> query(
      const T &x1, const T &y1, const T &x2, const T &y2
  ) const {
    assert(!(x2 < x1) && !(y2 < y1));
    std::vector<std::tuple<int, T, T>> res;
    query(0, 0, static_cast<int>(points.size()) - 1, x1, y1, x2, y2, res);
    return res;
  }
};

/*** 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}};
  RangeTree<int> t(v.begin(), v.end());
  auto got = t.query(-1, -1, 2, 5);
  sort(got.begin(), got.end());
  assert((got == vector<tuple<int, int, int>>{{0, 1, 4}, {2, 2, 2}, {8, -1, -1}, {9, 2, -1}}));
  got = t.query(1, 1, 4, 8);
  sort(got.begin(), got.end());
  assert((got == vector<tuple<int, int, int>>{{0, 1, 4}, {2, 2, 2}, {3, 3, 1}}));
  return 0;
}