HDU 4347 The Closest M Points【KD树】

题意：给你n个k维的点，再给你一个目标点x，让你查询离x最近的M个点？
分析：【KD树模板】

#include<cstdio>
#include<bits/stdc++.h>
#include<algorithm>
#include<string.h>
using namespace std;

const int N = 50000; //最大点数
short currentDim;
#define square(x) (x)*(x)
int n; //n个点
short K, t; //K维，t次询问

struct point {
    short pos[5];
    bool flag;
    point() {
        flag = false;
    }
    bool operator < (const point &a) const {
        return pos[currentDim] < a.pos[currentDim];
    }
};

point input_v[N];
priority_queue<pair<double, point>> pq;
point kd_tree[4*N];

inline void build(int n, short K, int p, int r, int index, short depth) {
    if(p <= r) {
        currentDim = depth % K;
        int mid = (p + r) / 2;
        nth_element(input_v+p, input_v+mid, input_v+r+1);
        kd_tree[index] = input_v[mid];
        kd_tree[index].flag = true;
        build(n, K, p, mid-1, 2*index, depth+1);
        build(n, K, mid+1, r, 2*index+1, depth+1);
    }
}

inline void query(point p, short closestM, short K, int index, short depth) {
    if(!kd_tree[index].flag) return;
    pair<double, point> current_node(0, kd_tree[index]);
    for(int i = 0; i < K; i++) {
        current_node.first += square(p.pos[i]-current_node.second.pos[i]);
    }
    short idm = depth % K;
    short flag = 0;
    int lchild = 2*index;
    int rchild = 2*index+1;
    if(kd_tree[index].pos[idm] <= p.pos[idm]) {
        query(p, closestM, K, rchild, depth+1);
    } else {
        query(p, closestM, K, lchild, depth+1);
    }
    if(pq.size() < closestM) {
        pq.push(current_node);
        flag = 1;
    } else {
        if(current_node.first < pq.top().first) {
            pq.pop();
            pq.push(current_node);
        }
        if(square(p.pos[idm] - kd_tree[index].pos[idm]) < pq.top().first) {
            flag = 1;
        }
    }
    if(flag) {
        if(kd_tree[index].pos[idm] <= p.pos[idm])
            query(p, closestM, K, lchild, depth+1);
        else
            query(p, closestM, K, rchild, depth+1);
    }
}

int main() {
    while(scanf("%d%hd",&n, &K) != EOF)   {
        for(int i = 0; i < n; i++) { //n个点
            for(short j = 0; j < K; j++) {
                scanf("%hd", &input_v[i].pos[j]);
            }
        }
        memset(kd_tree, 0, sizeof(point)*4*N);
        build(n, K, 0, n-1, 1, 0); //建树
        scanf("%hd", &t);
        for(short c = 0; c < t; c++) { //t次询问
            short closestM;
            point tmpv;
            for(short j = 0; j < K; j++) { //K维
                scanf("%hd", &tmpv.pos[j]);
            }
            scanf("%hd", &closestM); //最近M个点
            query(tmpv, closestM, K, 1, 0); //查询
            printf("the closest %hd points are:\n", closestM);
            point pt[10];
            for(short i = 0; !pq.empty(); i++) {
                pt[i] = pq.top().second; //最近M个点保存到pq队列
                pq.pop();
            }
            for(short i = closestM-1; i >= 0; i--) {
                for(short j = 0; j < K; j++) {
                    printf("%hd", pt[i].pos[j]);
                    if(j == K-1) printf("\n");
                    else printf(" ");
                }
            }
        }
    }
    return 0;
}