BZOJ 2626: JZPFAR KDtree + 堆

escription

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

Input

　　第一行，一个整数n，表示点的个数。
　　下面n行，每行两个整数x_i, y_i，表示n个点的坐标。点的标号按照输入顺序，分别为1..n。
　　下面一行，一个整数m，表示询问个数。
　　下面m行，每行三个整数px_i, py_i, k_i，表示一个询问。

Output

　　m行，每行一个整数，表示相应的询问的答案。

题解：  和之前一道题几乎一模一样，直接用堆 + $KDtree$ 维护即可.

#include<bits/stdc++.h>
#define maxn 200000
#define inf 1000000000000000
#define mid ((l+r)>>1)
#define rson(x) (t[x].ch[1])
#define lson(x) (t[x].ch[0])
#define ll long long
using namespace std;
void setIO(string s)
{
    string in = s + ".in";
    freopen(in.c_str(), "r" , stdin);
}
int d, n, m;
struct Node
{
    ll dis;
    int k;
    Node (ll dis = 0, int k = 0) : dis(dis), k(k) {}
    bool operator < (Node b) const
    {
        return dis == b.dis ? b.k > k : b.dis < dis;
    }
};
priority_queue <Node> Q;
// return if a is less than b
bool cmp1(Node a, Node b)
{
    return a.dis == b.dis ? a.k > b.k : a.dis < b.dis;
}
ll sqr(ll a)
{
    return a * a;
}
struct ND
{
    int ch[2], id;
    ll p[2], minv[2], maxv[2];
}t[maxn],T;
bool cmp(ND a, ND b)
{
    return a.p[d] == b.p[d] ? a.p[d ^ 1] < b.p[d ^ 1] : a.p[d] < b.p[d];
}
void pushup(int x, int y)
{
    for(int i = 0; i < 2 ; ++i)
    {
        t[x].minv[i] = min(t[x].minv[i], t[y].minv[i]);
        t[x].maxv[i] = max(t[x].maxv[i], t[y].maxv[i]);
    }
}
int build(int l, int r, int o)
{
    d = o;
    nth_element(t + l, t + mid, t + 1 + r, cmp);
    for(int i = 0; i < 2 ; ++i)
        t[mid].minv[i] = t[mid].maxv[i] = t[mid].p[i];
    lson(mid) = rson(mid) = 0;
    if(mid > l)
    {
        lson(mid) = build(l, mid - 1, o ^ 1);
        pushup(mid, lson(mid));
    }
    if(r > mid)
    {
        rson(mid) = build(mid + 1, r, o ^ 1);
        pushup(mid, rson(mid));
    }
    return mid;
}
ll getmax(int x)
{
    ll ans = 0;
    for(int i = 0; i < 2 ; ++i)
    {
        ans += max(sqr(t[x].minv[i] - T.p[i]), sqr(t[x].maxv[i] - T.p[i]));
    }
    return ans;
}
void query(int x, ll x1, ll y1)
{
    ll cur = sqr(t[x].p[0] - x1) + sqr(t[x].p[1] - y1), dl, dr, dn = getmax(x);
    if(dn < Q.top().dis) return;
    if(cmp1(Q.top(), Node(cur, t[x].id)))
    {
        Q.pop();
        Q.push(Node(cur, t[x].id));
    }
    dl = lson(x) ? getmax(lson(x)) : -inf;
    dr = rson(x) ? getmax(rson(x)) : -inf;
    if(dl > dr)
    {
        if(dl >= Q.top().dis) query(lson(x), x1, y1);
        if(dr >= Q.top().dis) query(rson(x), x1, y1);
    }
    else
    {
        if(dr >= Q.top().dis) query(rson(x), x1, y1);
        if(dl >= Q.top().dis) query(lson(x), x1, y1);
    }
}
int main()
{
    int i, j, k, root;
    ll x, y;
    // setIO("input");
    scanf("%d",&n);
    for(i = 1; i <= n ; ++i) scanf("%lld%lld",&t[i].p[0],&t[i].p[1]), t[i].id = i;
    root = build(1, n, 0);
    scanf("%d",&m);
    while(m --)
    {
        scanf("%lld%lld%d",&x,&y,&k);
        while(!Q.empty()) Q.pop();
        for(i = 1; i <= k ; ++i) Q.push(Node(-2333, -1));
        T.p[0] = x, T.p[1] = y;
        query(root, x, y);
        printf("%d\n",Q.top().k);
    }
    return 0;
}