BZOJ 2626 JZPFAR K-D樹

BZOJ 2626 JZPFAR K-D樹

題目大意：給出平面上的一些點，求到一個點的最遠的第k個點的標號。

思路：樸素的K-D樹建樹，然後在搜索的時候維護一個小跟堆，保留著最大的k個點，然後吧第k大的點作為基准點來判斷是否更新其他的點。

CODE：

#include 
#include 
#include 
#include 
#include 
#include 
#define MAX 100010
#define INF 0x3f3f3f3f 
using namespace std;

int dim;

struct Point{
	long long x,y;
	int id;
	
	Point(long long _ = 0,long long __ = 0):x(_),y(__) {}
	bool operator <(const Point &a)const {
		if(dim)	return x < a.x;
		return y < a.y;
	}
	void Read(int p) {
		scanf("%lld%lld",&x,&y);
		id = p;
	}
}point[MAX];

inline long long Calc(const Point &p1,const Point &p2)
{
	return (p1.x - p2.x) * (p1.x - p2.x) + (p1.y - p2.y) * (p1.y - p2.y);
}

struct KDTree{
	KDTree *son[2];
	Point root;
	long long x0,y0,x1,y1;
	
	KDTree(KDTree *_,KDTree *__,Point ___) {
		son[0] = _,son[1] = __;
		root = ___;
		x0 = x1 = ___.x;
		y0 = y1 = ___.y;
	}
	KDTree() {}
	void Maintain(KDTree *a) {
		x0 = min(x0,a->x0);
		x1 = max(x1,a->x1);
		y0 = min(y0,a->y0);
		y1 = max(y1,a->y1);
	}
	long long Dis(const Point &p) {
		long long re = 0;
		re = max(re,Calc(p,Point(x0,y0)));
		re = max(re,Calc(p,Point(x1,y0)));
		re = max(re,Calc(p,Point(x1,y1)));
		re = max(re,Calc(p,Point(x0,y1)));
		return re;
	}
}*root,none,*nil = &none;

struct Complex{
	long long dis;
	int id;
	
	Complex(long long _,int __):dis(_),id(__) {}
	bool operator <(const Complex &a)const {
		if(dis == a.dis)	return id < a.id;
		return dis > a.dis;
	}
};

int cnt,asks;

KDTree *BuildTree(int l,int r,int d)
{
	if(l > r)	return nil;
	dim = d;
	int mid = (l + r) >> 1;
	nth_element(point + l,point + mid,point + r + 1);
	KDTree *re = new KDTree(BuildTree(l,mid - 1,!d),BuildTree(mid + 1,r,!d),point[mid]);
	if(re->son[0] != nil)	re->Maintain(re->son[0]);
	if(re->son[1] != nil)	re->Maintain(re->son[1]);
	return re;
}

priority_queue q;

void Ask(KDTree *a,const Point &p)
{
	long long dis = Calc(p,a->root);
	Complex temp(dis,a->root.id);
	if(temp < q.top()) {
		q.push(temp);
		q.pop();
	}
	long long l = a->son[0] == nil ? -1:a->son[0]->Dis(p);
	long long r = a->son[1] == nil ? -1:a->son[1]->Dis(p);
	if(l > r) {
		if(a->son[0] != nil)
			Ask(a->son[0],p);
		if(a->son[1] != nil && r >= q.top().dis)
			Ask(a->son[1],p);
	}
	else {
		if(a->son[1] != nil)
			Ask(a->son[1],p);
		if(a->son[0] != nil && l >= q.top().dis)
			Ask(a->son[0],p);
	}
}

int main()
{
	cin >> cnt;
	for(int i = 1; i <= cnt; ++i)
		point[i].Read(i);
	root = BuildTree(1,cnt,0);
	cin >> asks;
	for(int k,i = 1; i <= asks; ++i) {
		Point p;
		p.Read(0);
		scanf("%d",&k);
		while(!q.empty())	q.pop();
		for(int i = 1; i <= k; ++i)	q.push(Complex(-INF,0));
		Ask(root,p);
		printf("%d\n",q.top().id);
	}
	return 0;
}