[KD-tree] Luogu P2093 JZPFAR

题目描述

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

题解

• 利用KD-Tree算法，求二维平面内到一定点的第k远点
• 正常建树，由于题目要求输出第k远点的编号，我们需记录一下每个节点的id，用优先队列动态维护小根堆，保持k个最优点，最后输出堆顶元素的编号即可

代码

#include <queue>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
#define sqr(x) (x)*(x)
#define N 100010
#define ll long long
using namespace std;
ll n,m,X,Y,tot,D,root;
struct node
{
    ll dis,id;
    bool operator <  (const node &a)const { return dis>a.dis||(dis==a.dis&&id<a.id); }
};
struct Node
{
    ll x[2],id;
    bool operator < (const Node &a) const {return x[D]<a.x[D];}
}p[N];
struct kdtree { Node p; ll mx[2],mn[2],ls,rs,id; }t[N];
priority_queue<node>Q;
ll dis(kdtree t) { return sqr(t.p.x[0]-X)+sqr(t.p.x[1]-Y); }
ll mxdis(kdtree t) { return max(sqr((t.mn[0]-X)),sqr(t.mx[0]-X))+max(sqr(t.mn[1]-Y),sqr(t.mx[1]-Y)); }
void update(ll d)
{
    if (!d) return;
    ll l=t[d].ls,r=t[d].rs;
    if (l) t[d].mn[0]=min(t[d].mn[0],t[l].mn[0]),t[d].mx[0]=max(t[d].mx[0],t[l].mx[0]),t[d].mn[1]=min(t[d].mn[1],t[l].mn[1]),t[d].mx[1]=max(t[d].mx[1],t[l].mx[1]);
    if (r) t[d].mn[0]=min(t[d].mn[0],t[r].mn[0]),t[d].mx[0]=max(t[d].mx[0],t[r].mx[0]),t[d].mn[1]=min(t[d].mn[1],t[r].mn[1]),t[d].mx[1]=max(t[d].mx[1],t[r].mx[1]);
}
void build(ll &d,ll l,ll r,ll x)
{
    if (l>r) return;
    ll mid=l+r>>1;
    D=x,d=++tot,nth_element(p+l,p+mid,p+r+1);
    t[d].p=p[mid],t[d].id=t[d].p.id,t[d].mn[0]=t[d].mx[0]=t[d].p.x[0],t[d].mn[1]=t[d].mx[1]=t[d].p.x[1];
    build(t[d].ls,l,mid-1,x^1),build(t[d].rs,mid+1,r,x^1),update(d);
}
void query(ll d)
{
    if (!d) return;
    ll p=dis(t[d]),l=t[d].ls,r=t[d].rs,L,R;
    if (p>Q.top().dis||(p==Q.top().dis&&t[d].id<Q.top().id)) Q.pop(),Q.push((node){p,t[d].id});
    if (l) L=mxdis(t[l]); if (r) R=mxdis(t[r]);
    if (L>R)
    {
        if (L>=Q.top().dis) query(l);
        if (R>=Q.top().dis) query(r);
    }
    else 
    {
        if (R>=Q.top().dis) query(r);
        if (L>=Q.top().dis) query(l);
    }
}
int main()
{
    scanf("%lld",&n);
    for (ll i=1;i<=n;i++) scanf("%lld%lld",&p[i].x[0],&p[i].x[1]),p[i].id=i;
    build(root,1,n,0),scanf("%lld",&m);
    for (ll i=1,k;i<=m;i++) 
    {
        scanf("%lld%lld%lld",&X,&Y,&k);
        while (Q.size()) Q.pop();
        for (ll j=1;j<=k;j++) Q.push((node){-1,0});
        query(root),printf("%lld
",Q.top().id);
    }
}