HDU 4347 The Closest M Points （kdTree）

#include<bits/stdc++.h>
using namespace std;
const int maxn = 5e4+6,K = 5;

#define squ(x) ((x)*(x))

int k,n,idx;

struct Point
{
    int x[K];
    bool operator <(const Point& rhs) const{
        return x[idx] < rhs.x[idx];
    }
    void print() const {
        for(int i = 0; i < k-1; i++)
            printf("%d ",x[i]);
        printf("%d\n",x[k-1]);
    }
}P[maxn];

#define fi first
#define se second
typedef pair<double,Point> HeapNode;
priority_queue<HeapNode> q;

struct kdTree
{
    Point Node[maxn<<2];
    int son[maxn<<2];
    #define lch (rt<<1)
    #define rch (rt<<1|1)

    void build(int l,int r,int rt = 1,int dep = 0)
    {
        if(l>r) return;
        son[rt] = r-l;
        int x = lch,y = rch;
        son[x] = son[y] = -1;
        idx = dep%k;
        int mid = (l+r)>>1;
        nth_element(P+l,P+mid,P+r+1);
        Node[rt] = P[mid];
        build(l,mid-1,x,dep+1);
        build(mid+1,r,y,dep+1);
    }

    Point qp; int qm;
    void query(int rt = 1,int dep = 0)
    {
        if(!~son[rt]) return;
        HeapNode u(0,Node[rt]);
        for(int i = 0; i < k; i++)
            u.fi += squ(u.se.x[i]-qp.x[i]);

        int dim = dep%k, x = lch, y = rch;
        bool flag = false;
        if(qp.x[dim]>=Node[rt].x[dim]) swap(x,y);
        if(~son[x]) query(x,dep+1);
        if(q.size()<qm) q.push(u),flag = true;
        else {
            if(u.fi<q.top().fi) q.pop(),q.push(u);
            if(squ(qp.x[dim]-Node[rt].x[dim])<q.top().fi) flag = true;
        }
        if(flag&&~son[y]) query(y,dep+1);
    }
}kd;

const int M = 11;
Point ans[M];

int main()
{
   // freopen("in.txt","r",stdin);
    while(~scanf("%d%d",&n,&k)){
        for(int i = 0; i < n; i++)
        for(int j = 0; j < k; j++)
            scanf("%d",&P[i].x[j]);
        kd.build(0,n-1);
        int t; scanf("%d",&t);
        while(t--){
            for(int i = 0; i < k; i++) scanf("%d",&kd.qp.x[i]);
            scanf("%d",&kd.qm);
            kd.query();
            printf("the closest %d points are:\n",kd.qm);
            int top = 0;
            while(q.size()){
                ans[++top] = q.top().se;
                q.pop();
            }
            while(top){
                ans[top].print();
                top--;
            }
        }
    }
    return 0;
}