BZOJ 2626 JZPFAR（KD-tree）

题目链接：http://61.187.179.132/JudgeOnline/problem.php?id=2626

题意：平面上有n个点。现在有m次询问，每次给定一个点(px, py)和一个整数k，输出n个点中离(px, py)的距离第k大的点的标号。如果有两个(或多个)点距离(px, py)相同，那么认为标号较小的点距离较大。

思路：将n个点建立KD-tree。

struct node
{
i64 d[2],L,R,Max[2],Min[2];
int id;
};

node a[N<<3];
int cmpD;

int cmp(node a,node b)
{
if(a.d[cmpD]!=b.d[cmpD]) return a.d[cmpD]<b.d[cmpD];
return a.d[!cmpD]<b.d[!cmpD];
}

void pushUp(int t)
{
int i,j;
if(a[t].L)
{
j=a[t].L;
FOR0(i,2)
{
upMax(a[t].Max[i],a[j].Max[i]);
upMin(a[t].Min[i],a[j].Min[i]);
}
}
if(a[t].R)
{
j=a[t].R;
FOR0(i,2)
{
upMax(a[t].Max[i],a[j].Max[i]);
upMin(a[t].Min[i],a[j].Min[i]);
}
}
}

int build(int L,int R,int d)
{
int M=(L+R)>>1;
cmpD=d;
nth_element(a+L,a+M,a+R+1,cmp);
a[M].Max[0]=a[M].Min[0]=a[M].d[0];
a[M].Max[1]=a[M].Min[1]=a[M].d[1];
if(L!=M) a[M].L=build(L,M-1,!d);
if(M!=R) a[M].R=build(M+1,R,!d);
pushUp(M);
return M;
}

i64 dist(int p,i64 x,i64 y)
{
i64 dis=0,d1,d2;
d1=abs(a[p].Max[0]-x);
d2=abs(a[p].Min[0]-x);
if(d1>d2) dis+=sqr(d1);
else dis+=sqr(d2);
d1=abs(a[p].Max[1]-y);
d2=abs(a[p].Min[1]-y);
if(d1>d2) dis+=sqr(d1);
else dis+=sqr(d2);
return dis;
}

i64 X,Y,K;
i64 ans[25];
int id[25];

void query(int p)
{
if(!p) return;
i64 dl,dr,d0,t=K,i;
d0=sqr(a[p].d[0]-X)+sqr(a[p].d[1]-Y);
while(t&&(ans[t]<d0||ans[t]==d0&&id[t]>a[p].id)) t--;
if(t!=K)
{
t++;
for(i=K;i>=t+1;i--)
{
ans[i]=ans[i-1];
id[i]=id[i-1];
}
ans[t]=d0; id[t]=a[p].id;
}
if(a[p].L) dl=dist(a[p].L,X,Y); else dl=0;
if(a[p].R) dr=dist(a[p].R,X,Y); else dr=0;
if(dl>dr)
{
if(dl>=ans[K]) query(a[p].L);
if(dr>=ans[K]) query(a[p].R);
}
else
{
if(dr>=ans[K]) query(a[p].R);
if(dl>=ans[K]) query(a[p].L);
}
}

int n,m,root;

int main()
{
RD(n);
int i;
FOR1(i,n) RD(a[i].d[0],a[i].d[1]),a[i].id=i;
root=build(1,n,0);
RD(m);
while(m--)
{
RD(X,Y,K);
FOR1(i,20) ans[i]=0,id[i]=INF;
query(root);
PR(id[K]);
}
}