HDU 4347-KD-Tree(K临近)

      k-dimension tree，能够存储k维数据的一棵树。这里这个树其实是指BST。朴素的BST，像Treap、Splay这些，节点都是存储一维的信息，但是KD树就是它们的拓展，可以存储多维数据。一维的时候很容易做到，直接把数据大小当作key，小的放在左边大的放在右边即可。其实拓展到高维的时候也是类似的，只不过要考虑多个维度，最显而易见的想法就是对于每一层，我用不同的维度的数值作为key来进行比较。例如说，对于一个二维数据，我第一层用第一维，第二层用第二维，第三层第一维，第四层第二维……这样，一棵树可以比较均匀的建立。下图给出了一个建立好的KD树的案例。                          

        

#include<bits/stdc++.h>
#define sq(x) (x)*(x)
#define N (55555)
 
using namespace std;
 
int idx,k,n,m,q;
 
struct Node
{
    int x[5];
    bool operator < (const Node &u) const
    {
        return x[idx] < u.x[idx];
    }
} P[N];
 
typedef pair<double,Node> PDN;
priority_queue<PDN> que;
 
struct KD_Tree
{
    int sz[N<<2]; Node p[N<<2];
 
    void build(int i,int l,int r,int dep)
    {
        if (l>r) return;
        int mid=(l+r)>>1;
        idx=dep%k;sz[i]=r-l;
        sz[i<<1]=sz[i<<1|1]=-1;
        nth_element(P+l,P+mid,P+r+1);
        p[i]=P[mid];
        build(i<<1,l,mid-1,dep+1);
        build(i<<1|1,mid+1,r,dep+1);
    }
 
    void query(int i,int m,int dep,Node a)
    {
        if (sz[i]==-1) return;
        PDN tmp=PDN(0,p[i]);
        for(int j=0;j<k;j++)
            tmp.first+=sq(tmp.second.x[j]-a.x[j]);
        int lc=i<<1,rc=i<<1|1,dim=dep%k,flag=0;
        if (a.x[dim]>=p[i].x[dim]) swap(lc,rc);
        if (~sz[lc]) query(lc,m,dep+1,a);
        if (que.size()<m) que.push(tmp),flag=1;
        else
        {
            if (tmp.first<que.top().first) que.pop(),que.push(tmp);
            if (sq(a.x[dim]-p[i].x[dim])<que.top().first) flag=1;
        }
        if (~sz[rc]&&flag) query(rc,m,dep+1,a);
    }
} KDT;
 
int main()
{
    while(~scanf("%d%d",&n,&k))
    {
        for(int i=0;i<n;i++)
            for(int j=0;j<k;j++)
                scanf("%d",&P[i].x[j]);
        KDT.build(1,0,n-1,0);
        scanf("%d",&q);
        while(q--)
        {
            Node now;
            for(int i=0;i<k;i++)
                scanf("%d",&now.x[i]);
            scanf("%d",&m); int t=0;
            KDT.query(1,m,0,now); Node pp[21];
            for(;!que.empty();que.pop())
                pp[++t]=que.top().second;
            printf("the closest %d points are:\n",t);
            for(int i=m;i>0;i--)
            {
                printf("%d",pp[i].x[0]);
                for(int j=1;j<k;j++)
                    printf(" %d",pp[i].x[j]);
                puts("");
            }
        }
    }
 
    return 0;
 
}