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]);
    }
}