bzoj 3053: The Closest M Points【KD-tree】

多维KDtree板子

左右儿子的估价用mn~mx当区间，假设区间里的数都存在；k维轮着做割点
#include<iostream>

#include<cstdio>

#include<algorithm>

#include<cmath>

#include<queue>

#include<cstring>

using namespace std;

const int N=50005;

int n,k,m,rt,w,ans[15];

priority_queue<pair<int,int> >q;

struct qwe

{

	int a[5];

	int& operator [] (int x)

	{

		return a[x];

	}

	bool operator < (const qwe &b) const

	{

		return a[w]<b.a[w];

	}

}a[N],b;

struct KD

{

	int ls,rs;

	qwe d,mn,mx;

}t[N<<2];

int read()

{

	int r=0,f=1;

	char p=getchar();

	while(p>'9'||p<'0')

	{

		if(p=='-')

			f=-1;

		p=getchar();

	}

	while(p>='0'&&p<='9')

	{

		r=r*10+p-48;

		p=getchar();

	}

	return r*f;

}

void minn(int &x,int y)

{

	x>y?x=y:0;

}

void maxx(int &x,int y)

{

	x<y?x=y:0;

}

void ud(int ro)

{

	if(t[ro].ls)

		for(int i=0;i<k;i++)

			minn(t[ro].mn[i],t[t[ro].ls].mn[i]),maxx(t[ro].mx[i],t[t[ro].ls].mx[i]);

	if(t[ro].rs)

		for(int i=0;i<k;i++)

			minn(t[ro].mn[i],t[t[ro].rs].mn[i]),maxx(t[ro].mx[i],t[t[ro].rs].mx[i]);

}

int build(int l,int r,int f)

{

	if(l>r)

		return 0;

	w=f;

	int mid=(l+r)>>1;

	nth_element(a+l,a+mid,a+r+1);

	t[mid].mn=t[mid].mx=t[mid].d=a[mid];

	t[mid].ls=build(l,mid-1,(f+1)%k);

	t[mid].rs=build(mid+1,r,(f+1)%k);

	ud(mid);

	return mid;

}

int dis(qwe a,qwe b)

{

	int r=0;

	for(int i=0;i<k;i++)

		r+=(a[i]-b[i])*(a[i]-b[i]);

	return r;

}

int wk(int ro)

{

	int r=0;

	for(int i=0;i<k;i++)

	{

		if(b[i]<t[ro].mn[i])

			r+=(t[ro].mn[i]-b[i])*(t[ro].mn[i]-b[i]);

		if(b[i]>t[ro].mx[i])

			r+=(t[ro].mx[i]-b[i])*(t[ro].mx[i]-b[i]);

	}

	return r;

}

void ques(int ro,int f)

{

	if(!ro)

		return;

	int dm=dis(t[ro].d,b),dl=t[ro].ls?wk(t[ro].ls):1e9,dr=t[ro].rs?wk(t[ro].rs):1e9;//cerr<<"OK"<<dm<<endl;

	if(q.top().first>dm)

		q.pop(),q.push(make_pair(dm,ro));

	if(dl<dr)

	{

		if(dl<q.top().first)

			ques(t[ro].ls,(f+1)%k);

		if(dr<q.top().first)

			ques(t[ro].rs,(f+1)%k);

	}

	else

	{

		if(dr<q.top().first)

			ques(t[ro].rs,(f+1)%k);

		if(dl<q.top().first)

			ques(t[ro].ls,(f+1)%k);

	}

}

int main()

{

	while(~scanf("%d%d",&n,&k))

	{

		memset(t,0,sizeof(t));

		for(int i=1;i<=n;i++)

			for(int j=0;j<k;j++)

				a[i][j]=read();

		rt=build(1,n,0);

		m=read();

		while(m--)

		{

			for(int i=0;i<k;i++)

				b[i]=read();

			int s=read();

			for(int i=1;i<=s;i++)

				q.push(make_pair(1e9,0));

			ques(rt,0);

			for(int i=1;i<=s;i++)

				ans[i]=q.top().second,q.pop();

			printf("the closest %d points are:\n",s);

			for(int i=s;i>=1;i--)

			{

				for(int j=0;j<k;j++)

					printf("%d ",t[ans[i]].d[j]);

				puts("");

			}

		}

	}

	return 0;

}
巴特西

bzoj 3053: The Closest M Points【KD-tree】

最新文章

热门文章
