hdu4347The Closest M Points KD树模板

Posted 2021-01-16 lqllulu

题意

一个k维空间，给出n个点的坐标，给出t个询问，每个询问给出一个点的坐标和一个m。对于每个询问找出跟这个点最接近的m个点

分析

 kd树的模板题。

  1 #include <cstdio>
  2 #include <cstring>
  3 #include <iostream>
  4 #include <algorithm>
  5 #include <queue>
  6 
  7 using namespace std;
  8 typedef long long LL;
  9 const int maxn=50080;
 10 const int K=5;
 11 int num,nownum,m;
 12 LL ans;
 13 struct kdNode{
 14     LL x[K];
 15     int div;
 16     bool lef;
 17 }Ans[12];
 18 struct Node{
 19     kdNode a;
 20     LL dis;
 21     bool operator <(const Node &a)const{
 22         return dis<a.dis;
 23     }
 24     Node(){}
 25     Node(kdNode &tmp,LL d){
 26         a=tmp;
 27         dis=d;
 28     }
 29 };
 30 int cmpNo;
 31 bool cmp(kdNode a,kdNode b){
 32     return a.x[cmpNo]<b.x[cmpNo];
 33 }
 34 inline LL max(LL a,LL b){//why?
 35     return a>b?a:b;
 36 }
 37 kdNode p[maxn],q;
 38 LL dis(kdNode a,kdNode b,int k){
 39     LL res=0;
 40     for(int i=0;i<k;i++){
 41         res+=(a.x[i]-b.x[i])*(a.x[i]-b.x[i]);
 42     }
 43     return res;
 44 }
 45 priority_queue<Node>qq;
 46 void buildKD(int l,int r,kdNode* p,int d,int k){
 47     if(l>r)return ;
 48     int m=(l+r)/2;
 49     cmpNo=d;
 50     nth_element(p+l,p+m,p+r+1,cmp);
 51     p[m].div=d;
 52     if(l==r){
 53         p[m].lef=1;
 54         return ;
 55     }
 56     buildKD(l,m-1,p,(d+1)%k,k);
 57     buildKD(m+1,r,p,(d+1)%k,k);
 58 }
 59 void findkd(int l,int r,kdNode& tar,kdNode* p,int k){
 60     if(l>r)return;
 61     int m=(l+r)/2;
 62     LL d=dis(p[m],tar,k);
 63     if(p[m].lef){
 64         if(nownum<num){
 65             nownum++;
 66             ans=max(ans,d);
 67             qq.push(Node(p[m],d));
 68         }
 69         else if(ans>d){
 70             qq.pop();
 71             qq.push(Node(p[m],d));
 72             ans=qq.top().dis;
 73         }
 74         return;
 75     }
 76     LL t=tar.x[p[m].div]-p[m].x[p[m].div];
 77     if(t>0){
 78         findkd(m+1,r,tar,p,k);
 79         if(nownum<num){
 80             qq.push(Node(p[m],d));
 81             nownum++;
 82             ans=qq.top().dis;
 83             findkd(l,m-1,tar,p,k);
 84         }else{
 85             if(ans>d){
 86                 qq.pop();
 87                 qq.push(Node(p[m],d));
 88                 ans=qq.top().dis;
 89             }
 90             if(ans>t*t)
 91                 findkd(l,m-1,tar,p,k);
 92         }
 93     }else{
 94         findkd(l,m-1,tar,p,k);
 95         if(nownum<num){
 96             qq.push(Node(p[m],d));
 97             nownum++;
 98             ans=qq.top().dis;
 99             findkd(m+1,r,tar,p,k);
100         }else{
101             if(ans>d){
102                 qq.pop();
103                 qq.push(Node(p[m],d));
104                 ans=qq.top().dis;
105             }
106             if(ans>t*t){
107                 findkd(m+1,r,tar,p,k);
108             }
109         }
110     }
111 }
112 
113 int n,k;
114 int main(){
115     while(scanf("%d%d",&n,&k)==2){
116         for(int i=0;i<n;i++){
117             for(int j=0;j<k;j++){
118                 scanf("%lld",&p[i].x[j]);
119             }
120             p[i].lef=0;
121         }
122         buildKD(0,n-1,p,k-1,k);
123         int t;
124         scanf("%d",&t);
125         for(int i=1;i<=t;i++){
126             ans=-1;
127             nownum=0;
128             for(int j=0;j<k;j++){
129                 scanf("%lld",&q.x[j]);
130             }
131             while(!qq.empty())qq.pop();
132             scanf("%d",&num);
133             findkd(0,n-1,q,p,k);
134             for(int j=1;j<=num;j++){
135                 Ans[j]=qq.top().a;
136                 qq.pop();
137             }
138             printf("the closest %d points are:
",num);
139             for(int j=num;j>=1;j--){
140                 for(int kk=0;kk<k;kk++){
141                     if(kk==0)
142                         printf("%lld",Ans[j].x[kk]);
143                     else
144                         printf(" %lld",Ans[j].x[kk]);
145                 }
146                 printf("
");
147             }
148         }
149     }
150 return 0;
151 }

View Code