BZOJ 2626 JZPFAR（KD-tree）

 1 #include<algorithm>
 2 #include<cstdio>
 3 #include<cmath>
 4 #include<cstring>
 5 #include<iostream>
 6 #define ll long long
 7 struct Point{
 8     ll d[2],max[2],min[2];
 9     int l,r,id;
10 }a[500005];
11 int root,n,K,cmpd,id[50],m;
12 ll X,Y,ans[50];
13 bool cmp(Point p1,Point p2){
14     if (p1.d[cmpd]!=p2.d[cmpd]) return p1.d[cmpd]<p2.d[cmpd];
15     else return p1.d[!cmpd]<p2.d[!cmpd];
16 }
17 void updata(int x){
18     if (a[x].l){
19         int j=a[x].l;
20         for (int i=0;i<2;i++)
21          a[x].max[i]=std::max(a[x].max[i],a[j].max[i]),
22          a[x].min[i]=std::min(a[x].min[i],a[j].min[i]);
23     }
24     if (a[x].r){
25         int j=a[x].r;
26         for (int i=0;i<2;i++)
27          a[x].max[i]=std::max(a[x].max[i],a[j].max[i]),
28          a[x].min[i]=std::min(a[x].min[i],a[j].min[i]);
29     }
30 }
31 ll sqr(ll x){
32     return x*x;
33 }
34 ll dist(int p,ll x,ll y){
35     ll dis=0,d1,d2;
36     d1=std::abs(x-a[p].max[0]);
37     d2=std::abs(x-a[p].min[0]);
38     if (d1>d2) dis+=d1*d1;
39     else dis+=d2*d2;
40     d1=std::abs(y-a[p].max[1]);
41     d2=std::abs(y-a[p].min[1]);
42     if (d1>d2) dis+=d1*d1;
43     else dis+=d2*d2;
44     return dis;
45 }
46 void query(int p){
47     if (!p) return;
48     ll dl,dr,d0;
49     int t=K;
50     d0=sqr(std::abs(a[p].d[0]-X))+sqr(std::abs(a[p].d[1]-Y));
51     while (t&&(d0>ans[t]||(d0==ans[t]&&a[p].id<id[t]))) t--;
52     if (t!=K){
53      t++;
54      for (int i=K;i>=t+1;i--){
55       ans[i]=ans[i-1];
56       id[i]=id[i-1];
57       }
58      ans[t]=d0;id[t]=a[p].id; 
59     }
60     if (a[p].l) dl=dist(a[p].l,X,Y);else dl=0;
61     if (a[p].r) dr=dist(a[p].r,X,Y);else dr=0;
62     if (dl>dr){
63         if (dl>=ans[K]) query(a[p].l);
64         if (dr>=ans[K]) query(a[p].r);
65     }else{
66         if (dr>=ans[K]) query(a[p].r);
67         if (dl>=ans[K]) query(a[p].l);
68     }
69 }
70 int build(int l,int r,int d){
71     int mid=(l+r)/2;
72     cmpd=d;
73     std::nth_element(a+l,a+mid,a+1+r,cmp);
74     a[mid].max[0]=a[mid].min[0]=a[mid].d[0];
75     a[mid].max[1]=a[mid].min[1]=a[mid].d[1];
76     if (l!=mid) a[mid].l=build(l,mid-1,!d);
77     if (r!=mid) a[mid].r=build(mid+1,r,!d);
78     updata(mid);
79     return mid;
80 }
81 int main(){
82     scanf("%d",&n);
83     for (int i=1;i<=n;i++){
84      scanf("%lld%lld",&a[i].d[0],&a[i].d[1]);
85      a[i].id=i;
86     }
87     root=build(1,n,0);
88     scanf("%d",&m);
89     while (m--){
90         scanf("%lld%lld%d",&X,&Y,&K);
91         for (int i=0;i<=20;i++) id[i]=99999999,ans[i]=0;
92         query(root);
93         printf("%d\n",id[K]); 
94     }
95 }