hdu4347

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

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

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

#include<bits/stdc++.h>

using namespace std;
const int inf=1e4+5;
int key,root,n,m,q,k,mxd;
int sqr(int x)
{
    return x*x;
}

struct point
{
    int d[7];
    friend int dis(point a,point b)
    {
        int s=0;
        for (int i=0; i<k; i++)
            s+=sqr(a.d[i]-b.d[i]);
        return s;
    }
} po;

struct node
{
    point nw;
    int son[2],mi[7],mx[7];
    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[500010];
    void init()
    {

        a[0].son[0]=a[0].son[1]=0;
        for (int i=0; i<5; i++)
        {
            a[0].mx[i]=-inf;
            a[0].mi[i]=inf;
        }
    }
    void update(int x)
    {
        int l=a[x].son[0],r=a[x].son[1];
        for (int i=0; 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 0;
        int m=(l+r)>>1;
        key=cur; nth_element(a+l,a+m,a+r+1);
        a[m].son[0]=build(l,m-1,(cur+1)%k);
        a[m].son[1]=build(m+1,r,(cur+1)%k);
        update(m);
        return m;
    }
    int getmi(int x)
    {
        int s=0;
        for (int i=0; i<k; i++)
            s+=sqr(max(po.d[i]-a[x].mx[i],0)+max(a[x].mi[i]-po.d[i],0));
        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[0],r=a[q].son[1],dl=2147483647,dr=2147483647;
        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=1; i<=n; i++)
            for (int j=0; j<k; j++)
                scanf("%d",&kd.a[i].nw.d[j]);
        root=kd.build(1,n,0);
        scanf("%d",&q);
        while (q--)
        {
            for (int i=0; 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=0; i<k; i++)
                {
                    printf("%d",ans.d[i]);
                    if (i!=k-1) printf(" "); else puts("");
                }
            }
        }
    }
}