K-D Tree

k-D Tree (KDT, k-Dimension Tree) is a data structure that can efficiently process $k$-dimensional spatial information.

When the number of nodes $n$ is much greater than $2^k$, a k-D Tree is highly efficient.

In programming-contest problems, $k=2$ is typical. When analyzing time complexity on this page, we treat $k$ as a constant.

Building the tree

A k-D Tree has the shape of a binary search tree. Each node in the binary search tree corresponds to a point in $k$-dimensional space. The points in each subtree lie within a $k$-dimensional hyperrectangle, and every point in that hyperrectangle belongs to the subtree.

Suppose we know the coordinates of points in $k$-dimensional space, with $n$ distinct points in total. To build a k-D Tree from them, proceed as follows:

1. If the current hyperrectangle contains only one point, return that point.

2. Choose a dimension and divide the current hyperrectangle into two hyperrectangles along that dimension.

3. Choose a split point: select a point along the chosen dimension; points whose value in that dimension is smaller than it go into one hyperrectangle (the left subtree), and the rest go into the other (the right subtree).

4. Use the selected point as the root of this subtree, recursively build the left and right subtrees from the two hyperrectangles, and maintain subtree information.

For intuition, consider an example with $k=2$.

The k-D Tree constructed from it may have the following shape:

Each node's coordinates are those of the selected split point; the $x$ or $y$ beside a non-leaf node indicates the selected splitting dimension.

This complexity is not guaranteed. For steps $2,3$, we introduce two optimizations:

1. Cycle through the $k$ dimensions so that every dimension is split within any consecutive $k$ levels.
2. Each time a split point is chosen along a dimension, choose the median on that dimension, so that the two resulting subtrees are as nearly equal in size as possible.

With optimization $2$, the height of the constructed k-D Tree is at most $\log n+O(1)$.

The bottleneck in constructing a k-D Tree is quickly selecting the median along a dimension, placing values smaller than the median to its left and the rest to its right. Sorting by that dimension with the sort function each time takes $O(n\log^2 n)$. In fact, finding the median of $n$ elements and placing it in its correct sorted position can be done in $O(n)$.

We recall the idea of quicksort. Each time, choose a number, put smaller values to its left and larger values to its right so that it occupies its correct sorted position, then recursively sort both sides. The expected complexity is $O(n\log n)$. However, because a k-D Tree only requires the median to be in its correct sorted position, we need recurse into only one side containing the median. This can be shown to have expected complexity $O(n)$. The algorithm library provides a function nth_element() with the same behavior: to find the value between s[l] and s[r] that, after sorting according to cmp, would occupy position s[mid], while ensuring values to the left of s[mid] are smaller than s[mid] and values to the right are larger than s[mid], simply write nth_element(s+l,s+mid,s+r+1,cmp).

With this idea, the time complexity of constructing a k-D Tree is $O(n\log n)$.

Operations in high-dimensional space

When querying information about all points inside a high-dimensional rectangular region, record the maximum and minimum coordinate in each dimension among the points in each node's subtree. If the rectangle corresponding to the current subtree does not intersect the query rectangle, do not search that subtree further; if it is completely contained in the query rectangle, return the sum of the weights of all points in the subtree; otherwise, check whether the current point lies in the query rectangle, update the answer, and recursively search both subtrees.

int query(int p) {
  if (!p) return 0;
  bool flag{false};
  for (int k : {0, 1}) flag |= (!(l.x[k] <= t[p].L[k] && t[p].R[k] <= h.x[k]));
  if (!flag) return t[p].sum;
  for (int k : {0, 1})
    if (t[p].R[k] < l.x[k] || h.x[k] < t[p].L[k]) return 0;
  int ans{0};
  flag = false;
  for (int k : {0, 1}) flag |= (!(l.x[k] <= t[p].x[k] && t[p].x[k] <= h.x[k]));
  if (!flag) ans = t[p].v;
  return ans += query(t[p].l) + query(t[p].r);
}

Complexity analysis

First consider the two-dimensional case. When querying a rectangle $R$, classify the nodes of the k-D Tree into three types:

1. Disjoint from $R$.
2. Completely contained in $R$.
3. Partially contained in $R$.

Clearly, the complexity of one query is the number of type-3 points. A type-3 point's rectangle either completely contains $R$ or is disjoint from it; the former contributes only $O(h)=O(\log n)$ points. We now analyze the number of the latter.

First, we may shift every side of the rectangle by $\epsilon$ so that the query rectangle passes through none of the existing points. This clearly does not change the set of points covered by the query rectangle.

Observe that each rectangle corresponding to a disjoint type-3 point must be crossed by one side of $R$. Therefore, we only need to count how many rectangles each side of $R$ crosses, that is, the maximum number of point-associated rectangles crossed by any line segment.

Consider a node $u$ with four grandchildren, where reaching each grandchild involves splitting once in each of the two dimensions. Observe that splitting a rectangle into four subrectangles this way lets a line segment parallel to a coordinate axis pass through at most two regions. Thus, starting from $u$, a query can descend into at most two grandchildren that still contain type-3 points (if the segment lies exactly on a split boundary this need not hold, but shifting the query rectangle's boundary avoids that case).

When building the tree, each point is the median of its entire subtree in the current splitting dimension, so the subtree size is halved. Thus, if the subtree of $u$ has size $n$, we obtain the following recurrence:

By the Master Theorem, $T(n)=O(\sqrt{n})$.

Generalizing the recurrence to $k$ dimensions gives $T(n)=2^{k-1}T(n/2^k)+O(1)$, and hence $T(n)=O(n^{1-\frac1k})$ (treating $k$ as a constant).

Insertions/deletions

If this $k$-dimensional point set is dynamic, meaning that points may be inserted or deleted, the balance of the k-D Tree cannot be guaranteed. Because of how a k-D Tree is constructed, rotations are not supported, and random priorities like those in an FHQ Treap cannot guarantee the complexity either. There are two common maintenance methods.

Square-root rebuilding

On insertion, first store the point to be inserted and rebuild after every $B$ insertions.

For deletion, simply mark the point. If stricter guarantees are needed, maintain the number of deleted points in the tree and rebuild when it reaches $B$.

The amortized update complexity is $O(n\log n/B)$ and the query complexity is $O(B+n^{1-\frac1k})$. If the two operation counts are of the same order, the optimum is $B=O(\sqrt{n\log n})$ (update $O(\sqrt{n\log n})$, query $O(\sqrt{n\log n}+n^{1-\frac1k})$).

Binary grouping

Consider maintaining several k-D Trees whose sizes are powers of $2$, with the sum of their sizes equal to $n$.

On insertion, create a k-D Tree of size $1$, then repeatedly merge trees of equal size (flattening and rebuilding them directly). In an implementation, it is enough to rebuild only once.

The trees that need to be merged always have sizes starting at $2^0$ with consecutive exponents. The complexity is similar to binary addition: amortized $O(n\log^2 n)$, because rebuilding itself costs a factor of $\log$.

For a query, query each tree separately. The complexity is $O\left(\sum_{i\geq0} (\frac n{2^i})^{1-\frac1k}\right)=O(n^{1-\frac1k})$.

Example problems

The 20M memory limit rules out all segment-tree-over-tree structures, and mandatory online processing rules out CDQ divide and conquer; only a k-D Tree can be used.

The following is reference code for binary grouping.

#include <algorithm>
#include <iostream>
using namespace std;
constexpr int N(2e5), LG{18};

struct pt {
  int x[2];
  int v, sum;
  int l, r;
  int L[2], R[2];
} t[N + 5], l, h;

int rt[LG];
int b[N + 5], cnt;

void upd(int p) {
  t[p].sum = t[t[p].l].sum + t[t[p].r].sum + t[p].v;
  for (int k : {0, 1}) {
    t[p].L[k] = t[p].R[k] = t[p].x[k];
    if (t[p].l) {
      t[p].L[k] = min(t[p].L[k], t[t[p].l].L[k]);
      t[p].R[k] = max(t[p].R[k], t[t[p].l].R[k]);
    }
    if (t[p].r) {
      t[p].L[k] = min(t[p].L[k], t[t[p].r].L[k]);
      t[p].R[k] = max(t[p].R[k], t[t[p].r].R[k]);
    }
  }
}

int build(int l, int r, int dep = 0) {
  int p{(l + r) >> 1};
  nth_element(b + l, b + p, b + r + 1,
              [dep](int x, int y) { return t[x].x[dep] < t[y].x[dep]; });
  int x{b[p]};
  if (l < p) t[x].l = build(l, p - 1, dep ^ 1);
  if (p < r) t[x].r = build(p + 1, r, dep ^ 1);
  upd(x);
  return x;
}

void append(int &p) {
  if (!p) return;
  b[++cnt] = p;
  append(t[p].l);
  append(t[p].r);
  p = 0;
}

int query(int p) {
  if (!p) return 0;
  bool flag{false};
  for (int k : {0, 1}) flag |= (!(l.x[k] <= t[p].L[k] && t[p].R[k] <= h.x[k]));
  if (!flag) return t[p].sum;
  for (int k : {0, 1})
    if (t[p].R[k] < l.x[k] || h.x[k] < t[p].L[k]) return 0;
  int ans{0};
  flag = false;
  for (int k : {0, 1}) flag |= (!(l.x[k] <= t[p].x[k] && t[p].x[k] <= h.x[k]));
  if (!flag) ans = t[p].v;
  return ans += query(t[p].l) + query(t[p].r);
}

int main() {
  int n;
  cin >> n;
  int lst{0};
  n = 0;
  while (true) {
    int op;
    cin >> op;
    if (op == 1) {
      int x, y, A;
      cin >> x >> y >> A;
      x ^= lst;
      y ^= lst;
      A ^= lst;
      t[++n] = {{x, y}, A};
      b[cnt = 1] = n;
      for (int sz{0};; ++sz)
        if (!rt[sz]) {
          rt[sz] = build(1, cnt);
          break;
        } else
          append(rt[sz]);
    } else if (op == 2) {
      cin >> l.x[0] >> l.x[1] >> h.x[0] >> h.x[1];
      l.x[0] ^= lst;
      l.x[1] ^= lst;
      h.x[0] ^= lst;
      h.x[1] ^= lst;
      lst = 0;
      for (int i{0}; i < LG; ++i) lst += query(rt[i]);
      cout << lst << "\n";
    } else
      break;
  }
  return 0;
}

Neighborhood queries

First build a 2-D Tree for these $n$ points.

Enumerate each node and, for each node, find the closest point other than that node to obtain the answer. A brute-force traversal of every node in a 2-D Tree takes $O(n)$ time per query, so pruning is needed. We can maintain the minimum and maximum coordinate in each dimension among all points in a subtree. Suppose the closest-pair distance found so far is $ans$; if the minimum distance from the query point to the rectangle containing all points in the subtree is at least $ans$, then the subtree cannot contain the answer and is skipped.

Moreover, we can use a heuristic search: if both subtrees of a node may contain the answer, search first in the subtree closer to the query point. The minimum distance from the query point to the rectangle corresponding to a subtree is the evaluation function for this problem.

#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <cstring>
#include <iomanip>
#include <iostream>
using namespace std;
constexpr int MAXN = 200010;
int n, d[MAXN], lc[MAXN], rc[MAXN];
double ans = 2e18;

struct node {
  double x, y;
} s[MAXN];

double L[MAXN], R[MAXN], D[MAXN], U[MAXN];

double dist(int a, int b) {
  return (s[a].x - s[b].x) * (s[a].x - s[b].x) +
         (s[a].y - s[b].y) * (s[a].y - s[b].y);
}

bool cmp1(node a, node b) { return a.x < b.x; }

bool cmp2(node a, node b) { return a.y < b.y; }

void maintain(int x) {
  L[x] = R[x] = s[x].x;
  D[x] = U[x] = s[x].y;
  if (lc[x])
    L[x] = min(L[x], L[lc[x]]), R[x] = max(R[x], R[lc[x]]),
    D[x] = min(D[x], D[lc[x]]), U[x] = max(U[x], U[lc[x]]);
  if (rc[x])
    L[x] = min(L[x], L[rc[x]]), R[x] = max(R[x], R[rc[x]]),
    D[x] = min(D[x], D[rc[x]]), U[x] = max(U[x], U[rc[x]]);
}

int build(int l, int r) {
  if (l > r) return 0;
  if (l == r) {
    maintain(l);
    return l;
  }
  int mid = (l + r) >> 1;
  double avx = 0, avy = 0, vax = 0, vay = 0;  // average variance
  for (int i = l; i <= r; i++) avx += s[i].x, avy += s[i].y;
  avx /= (double)(r - l + 1);
  avy /= (double)(r - l + 1);
  for (int i = l; i <= r; i++)
    vax += (s[i].x - avx) * (s[i].x - avx),
        vay += (s[i].y - avy) * (s[i].y - avy);
  if (vax >= vay)
    d[mid] = 1, nth_element(s + l, s + mid, s + r + 1, cmp1);
  else
    d[mid] = 2, nth_element(s + l, s + mid, s + r + 1, cmp2);
  lc[mid] = build(l, mid - 1), rc[mid] = build(mid + 1, r);
  maintain(mid);
  return mid;
}

double f(int a, int b) {
  double ret = 0;
  if (L[b] > s[a].x) ret += (L[b] - s[a].x) * (L[b] - s[a].x);
  if (R[b] < s[a].x) ret += (s[a].x - R[b]) * (s[a].x - R[b]);
  if (D[b] > s[a].y) ret += (D[b] - s[a].y) * (D[b] - s[a].y);
  if (U[b] < s[a].y) ret += (s[a].y - U[b]) * (s[a].y - U[b]);
  return ret;
}

void query(int l, int r, int x) {
  if (l > r) return;
  int mid = (l + r) >> 1;
  if (mid != x) ans = min(ans, dist(x, mid));
  if (l == r) return;
  double distl = f(x, lc[mid]), distr = f(x, rc[mid]);
  if (distl < ans && distr < ans) {
    if (distl < distr) {
      query(l, mid - 1, x);
      if (distr < ans) query(mid + 1, r, x);
    } else {
      query(mid + 1, r, x);
      if (distl < ans) query(l, mid - 1, x);
    }
  } else {
    if (distl < ans) query(l, mid - 1, x);
    if (distr < ans) query(mid + 1, r, x);
  }
}

int main() {
  cin.tie(nullptr)->sync_with_stdio(false);
  cin >> n;
  for (int i = 1; i <= n; i++) cin >> s[i].x >> s[i].y;
  build(1, n);
  for (int i = 1; i <= n; i++) query(1, n, i);
  cout << fixed << setprecision(4) << sqrt(ans) << '\n';
  return 0;
}

Similar to the previous problem, we replace the closest pair with the $k$-th farthest pair and change the evaluation function to the maximum distance from the query point to the rectangle corresponding to a subtree. Maintain the distances of the current $k$ farthest pairs in a min-heap; if a newly found pair has distance greater than the heap top, pop the top and insert the new distance. Use the heap top for pruning as well.

Because the problem asks for unordered pairs, swapping the two endpoints still gives the same pair, so every ordered pair is counted twice; therefore, multiply the input $k$ by $2$.

#include <algorithm>
#include <cstring>
#include <iostream>
#include <queue>
using namespace std;
constexpr int MAXN = 100010;
long long n, k;
priority_queue<long long, vector<long long>, greater<long long>> q;

struct node {
  long long x, y;
} s[MAXN];

bool cmp1(node a, node b) { return a.x < b.x; }

bool cmp2(node a, node b) { return a.y < b.y; }

long long lc[MAXN], rc[MAXN], L[MAXN], R[MAXN], D[MAXN], U[MAXN];

void maintain(int x) {
  L[x] = R[x] = s[x].x;
  D[x] = U[x] = s[x].y;
  if (lc[x])
    L[x] = min(L[x], L[lc[x]]), R[x] = max(R[x], R[lc[x]]),
    D[x] = min(D[x], D[lc[x]]), U[x] = max(U[x], U[lc[x]]);
  if (rc[x])
    L[x] = min(L[x], L[rc[x]]), R[x] = max(R[x], R[rc[x]]),
    D[x] = min(D[x], D[rc[x]]), U[x] = max(U[x], U[rc[x]]);
}

int build(int l, int r) {
  if (l > r) return 0;
  int mid = (l + r) >> 1;
  double av1 = 0, av2 = 0, va1 = 0, va2 = 0;  // average variance
  for (int i = l; i <= r; i++) av1 += s[i].x, av2 += s[i].y;
  av1 /= (r - l + 1);
  av2 /= (r - l + 1);
  for (int i = l; i <= r; i++)
    va1 += (av1 - s[i].x) * (av1 - s[i].x),
        va2 += (av2 - s[i].y) * (av2 - s[i].y);
  if (va1 > va2)
    nth_element(s + l, s + mid, s + r + 1, cmp1);
  else
    nth_element(s + l, s + mid, s + r + 1, cmp2);
  lc[mid] = build(l, mid - 1);
  rc[mid] = build(mid + 1, r);
  maintain(mid);
  return mid;
}

long long sq(long long x) { return x * x; }

long long dist(int a, int b) {
  return max(sq(s[a].x - L[b]), sq(s[a].x - R[b])) +
         max(sq(s[a].y - D[b]), sq(s[a].y - U[b]));
}

void query(int l, int r, int x) {
  if (l > r) return;
  int mid = (l + r) >> 1;
  long long t = sq(s[mid].x - s[x].x) + sq(s[mid].y - s[x].y);
  if (t > q.top()) q.pop(), q.push(t);
  long long distl = dist(x, lc[mid]), distr = dist(x, rc[mid]);
  if (distl > q.top() && distr > q.top()) {
    if (distl > distr) {
      query(l, mid - 1, x);
      if (distr > q.top()) query(mid + 1, r, x);
    } else {
      query(mid + 1, r, x);
      if (distl > q.top()) query(l, mid - 1, x);
    }
  } else {
    if (distl > q.top()) query(l, mid - 1, x);
    if (distr > q.top()) query(mid + 1, r, x);
  }
}

int main() {
  cin >> n >> k;
  k *= 2;
  for (int i = 1; i <= k; i++) q.push(0);
  for (int i = 1; i <= n; i++) cin >> s[i].x >> s[i].y;
  build(1, n);
  for (int i = 1; i <= n; i++) query(1, n, i);
  cout << q.top() << endl;
  return 0;
}

Exercises

“SDOI2010” Hide and Seek

“Violet” Angel Doll/SJY Place Chess Pieces

“National Training Team” JZPFAR

“BOI2007” Mokia

Luogu P4475 Chocolate Kingdom

“CH Weak Province Hu Ce R2” TATT