求与询问点欧几里德距离前m小的点

其实就是在kdtree询问的时候用优先队列维护一下就好了

好久没写kdtree练一练,注意这道题是多测

 #include<bits/stdc++.h>

 using namespace std;
const int inf=1e4+;
int key,root,n,m,q,k,mxd;
int sqr(int x)
{
return x*x;
} struct point
{
int d[];
friend int dis(point a,point b)
{
int s=;
for (int i=; i<k; i++)
s+=sqr(a.d[i]-b.d[i]);
return s;
}
} po; struct node
{
point nw;
int son[],mi[],mx[];
friend bool operator <(node a,node b)
{
return a.nw.d[key]<b.nw.d[key];
}
}; struct li
{
point a; int l;
friend bool operator <(li a, li b)
{
return a.l<b.l;
}
} mx;
set<li> st; struct kdtree
{
node a[];
void init()
{ a[].son[]=a[].son[]=;
for (int i=; i<; i++)
{
a[].mx[i]=-inf;
a[].mi[i]=inf;
}
}
void update(int x)
{
int l=a[x].son[],r=a[x].son[];
for (int i=; i<k; i++)
{
a[x].mi[i]=min(a[x].nw.d[i],min(a[l].mi[i],a[r].mi[i]));
a[x].mx[i]=max(a[x].nw.d[i],max(a[l].mx[i],a[r].mx[i]));
}
}
int build(int l,int r,int cur)
{
if (l>r) return ;
int m=(l+r)>>;
key=cur; nth_element(a+l,a+m,a+r+);
a[m].son[]=build(l,m-,(cur+)%k);
a[m].son[]=build(m+,r,(cur+)%k);
update(m);
return m;
}
int getmi(int x)
{
int s=;
for (int i=; i<k; i++)
s+=sqr(max(po.d[i]-a[x].mx[i],)+max(a[x].mi[i]-po.d[i],));
return s;
}
void ask(int q)
{
if (!q) return;
int tmp=dis(a[q].nw,po);
st.insert((li){a[q].nw,tmp});
if (st.size()>m)
{
set<li>::iterator it=st.end(); it--;
st.erase(it);
}
mxd=(*st.rbegin()).l;
int l=a[q].son[],r=a[q].son[],dl=,dr=;
if (l) dl=getmi(l);
if (r) dr=getmi(r);
if (dl<dr)
{
if (dl<mxd||st.size()<m) ask(l);
if (dr<mxd||st.size()<m) ask(r);
}
else {
if (dr<mxd||st.size()<m) ask(r);
if (dl<mxd||st.size()<m) ask(l);
}
}
} kd; int main()
{
while (scanf("%d%d",&n,&k)!=EOF)
{
kd.init();
for (int i=; i<=n; i++)
for (int j=; j<k; j++)
scanf("%d",&kd.a[i].nw.d[j]);
root=kd.build(,n,);
scanf("%d",&q);
while (q--)
{
for (int i=; i<k; i++)
scanf("%d",&po.d[i]);
scanf("%d",&m);
st.clear();
kd.ask(root);
printf("the closest %d points are:\n",m);
for (set<li>::iterator it=st.begin(); it!=st.end(); it++)
{
point ans=(*it).a;
for (int i=; i<k; i++)
{
printf("%d",ans.d[i]);
if (i!=k-) printf(" "); else puts("");
}
}
}
}
}
05-11 19:35