[hdu4347]The Closest M Points(线段树形式kd-tree)

解题关键：kdtree模板题，距离某点最近的m个点。

#include<cstdio>
#include<cstring>
#include<algorithm>
#include<cstdlib>
#include<iostream>
#include<cmath>
#include<queue>
#define sq(x) (x)*(x)
using namespace std;
typedef long long ll;
const int N=;

int idx,k,n,m,q;
struct node{
    int x[];
    bool operator<(const node &u)const{
        return x[idx]<u.x[idx];
    }
}P[N];
//线段树形式kd-tree
typedef pair<double,node>PDN;
priority_queue<PDN>que;
struct KD_Tree{
    int sz[N<<];
    node p[N<<];
    void build(int rt,int l,int r,int dep){
        if(l>r) return;
        int mid=(l+r)>>;
        idx=dep%k;sz[rt]=r-l;
        sz[rt<<]=sz[rt<<|]=-;
        nth_element(P+l,P+mid,P+r+);
        p[rt]=P[mid];
        build(rt<<,l,mid-,dep+);
        build(rt<<|,mid+,r,dep+);
    }

    void query(int i,int m,int dep,node a){//寻找m个。
        if(sz[i]==-) return;
        PDN tmp=PDN(,p[i]);//新建的tmp，first更新为0了
        for(int j=;j<k;j++)
            tmp.first+=sq(tmp.second.x[j]-a.x[j]);//距离目标点的距离
        int lc=i<<,rc=i<<|,dim=dep%k,flag=;
        if(a.x[dim]>=p[i].x[dim]) swap(lc,rc);
        if(~sz[lc]) query(lc,m,dep+,a);//~sz,sz!=-1
        if(que.size()<m) que.push(tmp),flag=;
        else{
            if(tmp.first<que.top().first) que.pop(),que.push(tmp);
            if(sq(a.x[dim]-p[i].x[dim])<que.top().first) flag=;//求到面的距离，空间相交
        }
        if(~sz[rc]&&flag) query(rc,m,dep+,a);
    }
}KDT;

int main(){
    while(~scanf("%d%d",&n,&k)){
        for(int i=;i<n;i++)
            for(int j=;j<k;j++)
                scanf("%d",&P[i].x[j]);
        KDT.build(,,n-,);
        scanf("%d",&q);
        while(q--){
            node now;
            for(int i=;i<k;i++)
                scanf("%d",&now.x[i]);
            scanf("%d",&m);
            KDT.query(,m,,now);
            node pp[];
            int t=;
            while(!que.empty()){
                pp[++t]=que.top().second;
                que.pop();
            }
            printf("the closest %d points are:\n",m);
            for(int i=m;i>;i--) for(int j=;j<k;j++) printf("%d%c",pp[i].x[j],j==k-?'\n':' ');
        }
    }
    return ;
}