KD树维护平面点集求近邻

一道KD树模板题：HDU4347 The Closest M Points。

AC代码：

#include 
using namespace std;

const int MAXN = 50000;
const int MAXK = 5;

int n, m, K, cmpIdx;
struct Coordinate {
    int d[MAXK];
    friend bool operator < (const Coordinate &a, const Coordinate &b) {
        return a.d[cmpIdx] < b.d[cmpIdx];
    }
} o, point[MAXN];

struct Node {
    int ls, rs;
    Coordinate coordinate;
} kdTree[MAXN];
int nodeCnt;

int calc_distance(const Coordinate &a, const Coordinate &b) {
    int d = 0;
    for (int i=0; i, distance_less > Q; // 与o的距离的大根堆
stack  S;

void push(const Coordinate &p) {
    Q.push(p);
    if (Q.size() > m) Q.pop();
} // 仅保留与o点距离较小的m个点

int build_kdTree(int l, int r, int idx) {
    if (l >= r) return 0;
    int m = (l + r) >> 1;
    cmpIdx = idx;
    nth_element(point + l, point + m, point + r); // 可以忽略坐标值相等的情况

    int ret = nodeCnt++;
    kdTree[ret].coordinate = point[m];
    kdTree[ret].ls = build_kdTree(l, m, (idx + 1) % K);
    kdTree[ret].rs = build_kdTree(m + 1, r, (idx + 1) % K);
    return ret;
}

void query_nn(int u, int idx) {
    cmpIdx = idx;
    int next_u = kdTree[u].ls, next_v = kdTree[u].rs;
    if (!(o < kdTree[u].coordinate)) swap(next_u, next_v);
    if (next_u) query_nn(next_u, (idx + 1) % K);

    /* 回溯 */
    push(kdTree[u].coordinate);
    int d = o.d[idx] - kdTree[u].coordinate.d[idx];
    if (next_v && calc_distance(Q.top(), o) >= d * d)
        query_nn(next_v, (idx + 1) % K);
}

int main() {
    while (scanf("%d %d", &n, &K) != EOF) {
        for (int i=0; i