【bzoj2626】JZPFAR KD-tree+堆

题目描述

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

输入

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

输出

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

样例输入

3
0 0
0 1
0 2
3
1 1 2
0 0 3
0 1 1

样例输出

3
1
1

---

题解

KD-tree+堆

求与一个点距离第k大的方法：维护一个k个元素的小根堆，表示距离，初始都为极小值；然后对于每个点，如果距离比堆顶元素大，则将堆顶元素替换为该距离；并根据左右子树的估价函数与堆顶元素的大小关系决定是否向下查询。

然而本题中要求距离相同的编号较小者距离较大，所以我们在以距离为第一关键字维护小根堆的同时以编号第二关键字维护大根堆即可。

注意在估价与堆顶相等时也要向下查询，因为可能存在编号的问题。

#include <queue>
#include <cstdio>
#include <cstring>
#include <utility>
#include <algorithm>
#define squ(x) (ll)(x) * (x)
#define N 100010
using namespace std;
typedef long long ll;
priority_queue<pair<ll , int> > q;
struct data
{
	int p[2] , mx[2] , mn[2] , id , c[2];
}a[N];
int d , root , num;
bool cmp(data a , data 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 k , int x)
{
	a[k].mx[0] = max(a[k].mx[0] , a[x].mx[0]);
	a[k].mn[0] = min(a[k].mn[0] , a[x].mn[0]);
	a[k].mx[1] = max(a[k].mx[1] , a[x].mx[1]);
	a[k].mn[1] = min(a[k].mn[1] , a[x].mn[1]);
}
int build(int l , int r , int now)
{
	int mid = (l + r) >> 1;
	d = now , nth_element(a + l , a + mid , a + r + 1 , cmp);
	a[mid].mx[0] = a[mid].mn[0] = a[mid].p[0];
	a[mid].mx[1] = a[mid].mn[1] = a[mid].p[1];
	if(l < mid) a[mid].c[0] = build(l , mid - 1 , now ^ 1) , pushup(mid , a[mid].c[0]);
	if(r > mid) a[mid].c[1] = build(mid + 1 , r , now ^ 1) , pushup(mid , a[mid].c[1]);
	return mid;
}
ll getdis(int k , int x , int y)
{
	return max(squ(a[k].mx[0] - x) , squ(a[k].mn[0] - x)) + max(squ(a[k].mx[1] - y) , squ(a[k].mn[1] - y));
}
void query(int k , int x , int y)
{
	ll dn = squ(a[k].p[0] - x) + squ(a[k].p[1] - y) , dl = (a[k].c[0] ? getdis(a[k].c[0] , x , y) : -1) , dr = (a[k].c[1] ? getdis(a[k].c[1] , x , y) : -1);
	if(make_pair(-dn , a[k].id) < q.top())
		q.pop() , q.push(make_pair(-dn , a[k].id));
	if(dl > dr)
	{
		if(dl >= -q.top().first) query(a[k].c[0] , x , y);
		if(dr >= -q.top().first) query(a[k].c[1] , x , y);
	}
	else
	{
		if(dr >= -q.top().first) query(a[k].c[1] , x , y);
		if(dl >= -q.top().first) query(a[k].c[0] , x , y);
	}
}
int main()
{
	int n , i , m , x , y , k;
	scanf("%d" , &n);
	for(i = 1 ; i <= n ; i ++ ) scanf("%d%d" , &a[i].p[0] , &a[i].p[1]) , a[i].id = i;
	root = build(1 , n , 0);
	scanf("%d" , &m);
	while(m -- )
	{
		scanf("%d%d%d" , &x , &y , &k);
		while(!q.empty()) q.pop();
		for(i = 1 ; i <= k ; i ++ ) q.push(make_pair(0 , 1 << 30));
		query(root , x , y);
		printf("%d\n" , q.top().second);
	}
	return 0;
}