[hdu4347]The Closest M Points(平衡树式kdtree)

解题关键：模板题（结合起来了）

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<queue>
#include<cmath>
#include<cstdlib>
#define N 50050
#define inf (1<<30)
#define sq(x) (x)*(x)
using namespace std;
int n,m,dim,rt,k;
struct node{
    int p[10],minn[10],maxx[10];
    bool operator<(const node &u)const{
        return p[dim]<u.p[dim];
    }
}a[N];
typedef pair<double,node>PDN;
priority_queue<PDN>que;
struct kd_tree{
    int c[N][2];
    node s[N],q;
    void update(int o){//管辖范围
        int l=c[o][0],r=c[o][1];
        for(int i=0;i<k;i++){
            if(l){ s[o].minn[i]=min(s[o].minn[i],s[l].minn[i]); s[o].maxx[i]=max(s[o].maxx[i],s[l].maxx[i]); }
            if(r){ s[o].minn[i]=min(s[o].minn[i],s[r].minn[i]); s[o].maxx[i]=max(s[o].maxx[i],s[r].maxx[i]); }
        }
    }
    void add(int o,node t){ for(int i=0;i<k;i++)s[o].minn[i]=s[o].maxx[i]=s[o].p[i]=t.p[i]; c[o][0]=c[o][1]=0;}
    int dist(node t,int o){
        int tmp=0;
        for(int i=0;i<k;i++){
            if(s[o].minn[i]>t.p[i])tmp+=sq(s[o].minn[i]-t.p[i]);
            if(t.p[i]>s[o].maxx[i])tmp+=sq(t.p[i]-s[o].maxx[i]);
        }
        return tmp;
    }
    void build(int &o,int l,int r,int now){
        o=(l+r)>>1;now%=k;dim=now;
        nth_element(a+l,a+o,a+r+1);
        add(o,a[o]);
        if(l!=o) build(c[o][0],l,o-1,now+1);
        if(o!=r) build(c[o][1],o+1,r,now+1);
        update(o);
    }
    void ins(int o,int now){
        now%=k;
        if(q.p[now]<s[o].p[now]){
            if(c[o][0]) ins(c[o][0],now+1);
            else c[o][0]=++n,add(n,q);
        }
        else{
            if(c[o][1]) ins(c[o][1],now+1);
            else c[o][1]=++n,add(n,q);
        }
        update(o);
    }
    void qry(int o){//曼哈顿距离,且只求最短，dis是最短距离
        PDN tmp=PDN(0,s[o]);
        for(int i=0;i<k;i++) tmp.first+=sq(s[o].p[i]-q.p[i]);
        if(que.size()<m) que.push(tmp);
        else{
            if(que.top().first>tmp.first){
                que.pop();
                que.push(tmp);
            }
        }
        int dl=c[o][0]?dist(q,c[o][0]):inf,dr=c[o][1]?dist(q,c[o][1]):inf;
        if(dl<dr){
            if((dl<que.top().first||que.size()<m)&&dl!=inf) qry(c[o][0]);
            if((dr<que.top().first||que.size()<m)&&dr!=inf) qry(c[o][1]);
        }else{
            if((dr<que.top().first||que.size()<m)&&dr!=inf) qry(c[o][1]);
            if((dl<que.top().first||que.size()<m)&&dl!=inf) qry(c[o][0]);
        }
    }
}kd;

int main(){
    while(scanf("%d%d",&n,&k)!=EOF){
        for(int i=1;i<=n;i++){
            for(int j=0;j<k;j++)
                scanf("%d",&a[i].p[j]);
        }
        kd.build(rt,1,n,0);
        int t;
        scanf("%d",&t);
        while(t--){
            while(!que.empty()) que.pop();
            for(int j=0;j<k;j++) scanf("%d",&kd.q.p[j]);
            scanf("%d",&m);
            kd.qry(rt);
            node pp[100];
            int nn=0;
            printf("the closest %d points are:\n",m);
            while(!que.empty()) pp[++nn]=que.top().second,que.pop();
            for(int i=m;i>0;i--){
                for(int j=0;j<k;j++){
                    printf("%d%c",pp[i].p[j],j==k-1?'\n':' ');
                }
            }
        }
    }
    return 0;
}