bzoj3053: The Closest M Points（k-dtree）

题目

#include<bits/stdc++.h>
using namespace std;
const int N=50002;
struct node{
	int l,r,d[5],mn[5],mx[5];
}t[N];
int n,k,i,j,Q[5],ans[12],now,rt,T,m;
priority_queue<pair<int,int> >q;
#define sqr(x) (x)*(x)
#define mp make_pair
#define gc getchar
inline int rd(){
	int x=0,fl=1;char ch=gc();
	for (;ch<48||ch>57;ch=gc())if(ch=='-')fl=-1;
	for (;48<=ch&&ch<=57;ch=gc())x=(x<<3)+(x<<1)+(ch^48);
	return x*fl;
}
inline void wri(int a){if(a<0)a=-a,putchar('-');if(a>=10)wri(a/10);putchar(a%10|48);}
inline void wln(int a){wri(a);puts("");}
int dis(int x){
	int ans=0;
	for (int i=0;i<k;i++) ans+=sqr(t[x].d[i]-Q[i]);
	return ans;
}
int get(int x){//x代表的块内离q可能的最近距离
	int ans=0;
	for (int i=0;i<k;i++){
		if (Q[i]<t[x].mn[i]) ans+=sqr(Q[i]-t[x].mn[i]);
		if (Q[i]>t[x].mx[i]) ans+=sqr(Q[i]-t[x].mx[i]);
	}
	return ans;
}
void up(int x){
	int l=t[x].l,r=t[x].r;
	for (int i=0;i<k;i++) t[x].mn[i]=min(t[x].mn[i],min(t[l].mn[i],t[r].mn[i])),t[x].mx[i]=max(t[x].mx[i],max(t[l].mx[i],t[r].mx[i]));
}
bool cmp(node a,node b){return a.d[now]<b.d[now];}
int build(int l,int r,int d){
	now=d;
	int mid=l+r>>1;
	nth_element(t+l,t+mid,t+r+1,cmp);
	if (l<mid) t[mid].l=build(l,mid-1,(d+1)%k);
	else t[mid].l=0;
	if (mid<r) t[mid].r=build(mid+1,r,(d+1)%k);
	else t[mid].r=0;
	up(mid);
	return mid;
}
void query(int x){
	if (!x) return;
	int d=dis(x),dl=get(t[x].l),dr=get(t[x].r);
	if (d<q.top().first) q.pop(),q.push(mp(d,x));
	if (dl<dr){
		if (dl<q.top().first) query(t[x].l);
		if (dr<q.top().first) query(t[x].r);
	}else{
		if (dr<q.top().first) query(t[x].r);
		if (dl<q.top().first) query(t[x].l);
	}
}
int main(){
while (~scanf("%d%d",&n,&k)){
	for (i=0;i<k;i++) t[0].mn[i]=1e9,t[0].mx[i]=-1e9;
	for (i=1;i<=n;i++)
		for (j=0;j<k;j++) t[i].mn[j]=t[i].mx[j]=t[i].d[j]=rd();
	rt=build(1,n,0);
	for (T=rd();T--;){
		for (i=0;i<k;i++) Q[i]=rd();
		m=rd(),printf("the closest %d points are:\n",m);
		for (i=1;i<=m;i++) q.push(mp(1e9,0));
		query(rt);
		for (i=1;i<=m;i++) ans[i]=q.top().second,q.pop();
		for (i=m;i;i--)
			for (j=0;j<k;j++) wri(t[ans[i]].d[j]),putchar(j==k-1?'\n':' ');
	}
}
}