hdu 4347 The Closest M Points (kd树)

Posted 2020-08-12

hdu 4347

题意：

　　求k维空间中离所给点最近的m个点，并按顺序输出  。

解法：

　　kd树模板题 。

　　不懂kd树的可以先看看这个 。

　　不多说，上代码 。

  1 #include <iostream>
  2 #include <cstdio>
  3 #include <algorithm>
  4 #include <cmath>
  5 #include <cstring>
  6 #include <queue>
  7 #include <set>
  8 #include <vector>
  9 #include <map>
 10 #define ll long long
 11 
 12 using namespace std;
 13 
 14 const int N=50007;
 15 const int K=6;
 16 
 17 int n,m;
 18 
 19 struct point{
 20     int a[K];
 21     int div;  // 按哪个维度划分
 22     bool lef;  // 是否是叶子节点
 23     ll dis;  // 离询问点的距离。注意这个在读入建树时不会用到，在进入队列时才用到
 24     void print(){
 25         printf("%d",a[0]);
 26         for (int i=1;i<m;i++) 
 27             printf(" %d",a[i]);
 28         puts("");
 29     }
 30     bool operator < (const point &t) const{
 31         return dis<t.dis;
 32     }
 33     point(){}
 34     point(point &t,ll d){
 35         for (int i=0;i<m;i++) a[i]=t.a[i];
 36         dis=d;
 37     }
 38 }p[N],tar;
 39 
 40 int cmp_NO;
 41 inline bool cmp(point x,point y){
 42     return x.a[cmp_NO]<y.a[cmp_NO];
 43 }
 44 
 45 inline ll dis(point x,point y){
 46     ll ret=0;
 47     for (int i=0;i<m;i++) 
 48         ret+=(x.a[i]-y.a[i])*(x.a[i]-y.a[i]);
 49     return ret;
 50 }
 51 
 52 inline void bulid_kdt(int L,int R,int d){
 53     if (L>R) return;
 54     int mid=(L+R)>>1;
 55     cmp_NO=d;
 56     nth_element(p+L,p+mid,p+R+1,cmp);
 57     p[mid].div=d;
 58     if (L==R){
 59         p[L].lef=true;
 60         return;
 61     }
 62     bulid_kdt(L,mid-1,(d+1)%m);
 63     bulid_kdt(mid+1,R,(d+1)%m);
 64 }
 65 
 66 priority_queue<point> que;
 67 int num,nownum;
 68 ll ans;
 69 
 70 inline void find_kd(int L,int R){
 71     if (L>R) return;
 72     
 73     int mid=(L+R)>>1;
 74     ll d=dis(p[mid],tar);
 75     if (p[mid].lef){
 76         if (nownum<num){
 77             nownum++;
 78             que.push(point(p[mid],d));
 79             ans=max(ans,d);
 80         }
 81         else if (ans>d){
 82             que.pop();
 83             que.push(point(p[mid],d));
 84             ans=que.top().dis;
 85         }
 86         return;
 87     }
 88 
 89     int t=tar.a[p[mid].div]-p[mid].a[p[mid].div];
 90     if (t>0){
 91         find_kd(mid+1,R);
 92         if (nownum<num){
 93             nownum++;
 94             que.push(point(p[mid],d));
 95             ans=max(ans,d);
 96             find_kd(L,mid-1);
 97         }
 98         else {
 99             if (ans>d){
100                 que.pop();
101                 que.push(point(p[mid],d));
102                 ans=que.top().dis;
103             }
104             if (ans>t*t) 
105                 find_kd(L,mid-1);
106         }
107     }
108     else {
109         find_kd(L,mid-1);
110         if (nownum<num){
111             nownum++;
112             que.push(point(p[mid],d));
113             ans=max(ans,d);
114             find_kd(mid+1,R);
115         }
116         else{
117             if (ans>d){
118                 que.pop();
119                 que.push(point(p[mid],d));
120                 ans=que.top().dis;
121             }
122             if (ans>t*t) 
123                 find_kd(mid+1,R);
124         }
125     }
126 }
127 
128 inline void put(){
129     if (que.empty()) return;
130     point pp=que.top();
131     que.pop();
132     put();
133     pp.print();
134 }
135 
136 int main(){
137     while (~scanf("%d%d",&n,&m)){
138         for (int i=0;i<n;i++){
139             for (int j=0;j<m;j++) 
140                 scanf("%d",&p[i].a[j]);
141             p[i].lef=false;
142         }
143 
144         bulid_kdt(0,n-1,0);  // 这一步相当于将原数组重新排了个序，先访问到的点放在中间
145 
146         int q;
147         scanf("%d",&q);
148         while (q--){
149             for (int i=0;i<m;i++)
150                 scanf("%d",&tar.a[i]);
151             while (!que.empty()) que.pop();
152             scanf("%d",&num);
153             nownum=0;
154             ans=-1;
155             find_kd(0,n-1);
156             printf("the closest %d points are:\n",num);
157             put();
158         }
159     }
160     return 0;
161 }