bzoj 2626: JZPFAR （KD-tree）

2626: JZPFAR

Time Limit: 50 Sec   Memory Limit: 128 MB
Submit: 961   Solved: 362
[ Submit][ Status][ Discuss]

Description

　　平面上有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行，每行一个整数，表示相应的询问的答案。

Sample Input

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

Sample Output

3
1
1

数据规模和约定
　　50%的数据中，n个点的坐标在某范围内随机分布。
　　100%的数据中，n<=10^5, m<=10^4, 1<=k<=20，所有点(包括询问的点)的坐标满足绝对值<=10^9，n个点中任意两点坐标不同，m个询问的点的坐标在某范围内随机分布。

HINT

Source

[ Submit][ Status][ Discuss] 

题解：KD-tree

#include<iostream>
#include<cstring>
#include<cstdio>
#include<algorithm>
#include<queue>
#define N 100003
#define LL long long
using namespace std;
int n,m,root,cmpd,x,y,k;
struct data {
	int l,r,num;
	int d[2],mx[2],mn[2];
	LL dis;
	bool operator <(const data &a) const
	{
		return dis>a.dis||dis==a.dis&&num<a.num;
	}
}tr[N];
priority_queue<data> q;
int cmp(data a,data b)
{
	return a.d[cmpd]<b.d[cmpd]||a.d[cmpd]==b.d[cmpd]&&a.d[cmpd^1]<b.d[cmpd^1];
}
void update(int now)
{
	int l=tr[now].l; int r=tr[now].r;
	for (int i=0;i<=1;i++) {
		if (l) tr[now].mx[i]=max(tr[now].mx[i],tr[l].mx[i]),
		       tr[now].mn[i]=min(tr[now].mn[i],tr[l].mn[i]);
		if (r) tr[now].mx[i]=max(tr[now].mx[i],tr[r].mx[i]),
		       tr[now].mn[i]=min(tr[now].mn[i],tr[r].mn[i]);
	}
}
int build(int l,int r,int d)
{
	cmpd=d;
	int mid=(l+r)/2;
	nth_element(tr+l,tr+mid,tr+r+1,cmp);
	//for (int i=l;i<=r;i++) 
	 //cout<<tr[i].d[0]<<" "<<tr[i].d[1]<<endl;
	//cout<<endl;
	for (int i=0;i<=1;i++)
	 tr[mid].mx[i]=tr[mid].mn[i]=tr[mid].d[i];
	if (l<mid) tr[mid].l=build(l,mid-1,d^1);
	if (r>mid) tr[mid].r=build(mid+1,r,d^1);
	update(mid);
	//cout<<tr[mid].d[0]<<" "<<tr[mid].d[1]<<endl;
	return mid;
}
LL pow(int x)
{
	return (LL)x*(LL)x;
}
LL dist(int now)
{
	LL t=0;
	t+=max(pow(tr[now].mx[0]-x),pow(tr[now].mn[0]-x));
	t+=max(pow(tr[now].mx[1]-y),pow(tr[now].mn[1]-y));
	return t;
}
void query(int now)
{
	LL d0,dl=0,dr=0;
	d0=pow(tr[now].d[0]-x)+pow(tr[now].d[1]-y);
	tr[now].dis=d0;
	if (q.size()<k) q.push(tr[now]);
	else {
		data t=q.top();
		if(t.dis<d0||t.dis==d0&&tr[now].num<t.num) {
			q.pop();
			q.push(tr[now]);
		}
	}
	if (tr[now].l) dl=dist(tr[now].l);
	if (tr[now].r) dr=dist(tr[now].r);
	if (dl>dr) {
		if ((q.size()<k||q.top().dis<=dl)&&tr[now].l) query(tr[now].l);
		if ((q.size()<k||q.top().dis<=dr)&&tr[now].r) query(tr[now].r);
	}
	else {
		if ((q.size()<k||q.top().dis<=dr)&&tr[now].r) query(tr[now].r);
		if ((q.size()<k||q.top().dis<=dl)&&tr[now].l) query(tr[now].l);
	}
}
int main()
{
	freopen("a.in","r",stdin);
	freopen("my.out","w",stdout);
	scanf("%d",&n);
	for (int i=1;i<=n;i++)
	 scanf("%d%d",&tr[i].d[0],&tr[i].d[1]),tr[i].num=i;
    root=build(1,n,0);
    scanf("%d",&m);
    for (int i=1;i<=m;i++) {
    	scanf("%d%d%d",&x,&y,&k);
    	while (!q.empty()) q.pop();
    	query(root);
        printf("%d\n",q.top().num);
	}
}