BZOJ 2626: JZPFAR KDtree + 堆

Code:

#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; 
}