[BZOJ 2626]JZPFAR

K-D tree的范围查找

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <queue>
#include <cmath>
#define maxn 100010
using namespace std;
typedef long long ll;
void read(int& num){
	num = 0;
	int f = 1;
	char ch = getchar();
	for(; ch < '!'; ch = getchar());
	if(ch == '-'){
		ch = getchar();
		f = -1;
	}
	for(; ch > '!'; ch = getchar())
		num = num * 10 + ch - 48;
	num *= f;
}

int n, q, D;
#define N 2
struct Point{
	int l, r, d[N], mn[N], mx[N], id;
	Point(int x = 0, int y = 0){
		r = l = 0;
		d[0] = x, d[1] = y;
	}
	bool operator<(const Point& k)const{
		return d[D] < k.d[D];
	}
	int& operator[](int i){return d[i];}
}t[maxn << 1], P;

int root;

void update(int x){
	int l = t[x].l, r = t[x].r;
	for(int i = 0; i < N; i ++){
		t[x].mn[i] = t[x].mx[i] = t[x][i];
		if(l){
			t[x].mn[i] = min(t[x].mn[i], t[l].mn[i]);
			t[x].mx[i] = max(t[x].mx[i], t[l].mx[i]);
		}
		if(r){
			t[x].mn[i] = min(t[x].mn[i], t[r].mn[i]);
			t[x].mx[i] = max(t[x].mx[i], t[r].mx[i]);
		}
	}
}

struct Pair{
	ll dis;
	int id;
	bool operator<(const Pair& k)const{
		if(dis != k.dis)return dis > k.dis;
		return id < k.id;
	}
	Pair(int x, ll y){
		id = x;
		dis = y;
	}
};

priority_queue<Pair>Q;


int build(int l, int r, int now){
	if(l > r)return 0;
	int mid = l + r >> 1;
	D = now;
	nth_element(t + l, t + mid, t + r + 1);
	t[mid].l = build(l, mid - 1, now ^ 1);
	t[mid].r = build(mid + 1, r, now ^ 1);
	update(mid);
	return mid;
}

const int inf = 0x7fffffff;

ll Sqr(ll x){
	return x * x;
}

ll dis(Point k){
	ll ret = 0;
	for(int i = 0; i < N; i ++)
	    ret += Sqr(k[i] - P[i]);
	return ret;
}

ll dist(Point k){
	ll ret = 0;
	for(int i = 0; i < N; i ++)
	    ret += max(Sqr(P[i] - k.mn[i]), Sqr(k.mx[i] - P[i]));
	return ret;
}

void ask(int root, int now){
	ll dl = -inf, dr = -inf;
	int l = t[root].l, r = t[root].r;
	ll d = dis(t[root]);
	if(d > Q.top().dis || (Q.top().dis == d && Q.top().id > t[root].id)){
		Q.pop();
		Q.push(Pair(t[root].id, d));
	}
	if(l)dl = dist(t[l]);
	if(r)dr = dist(t[r]);
	if(dl > dr){
		if(l && Q.top().dis <= dl)ask(l, now ^ 1);
		if(r && Q.top().dis <= dr)ask(r, now ^ 1);
	}
	else{
        if(r && Q.top().dis <= dr)ask(r, now ^ 1);
		if(l && Q.top().dis <= dl)ask(l, now ^ 1);
	}
}

int main(){
	read(n);
	int x, y, k;
	for(int i = 1; i <= n; i ++){
		read(x), read(y);
		t[i] = Point(x, y);
		t[i].id = i;
	}
	root = build(1, n, 0);
	read(q);
	while(q --){
		read(x), read(y), read(k);
		for(int i = 1; i <= k; i ++)
		    Q.push(Pair(0, -inf));
		P = Point(x, y);
		ask(root, 0);
		printf("%d\n", Q.top().id);
		while(!Q.empty())
			Q.pop();
	}
    return 0;
}