HDU 4347 The Closest M Points （kdTree）

赤果果的kdTree。

学习传送门：http://www.cnblogs.com/v-July-v/archive/2012/11/20/3125419.html

其实就是二叉树的变形

#include<bits/stdc++.h>

using namespace std;

const int maxn = 5e4+,K = ;

#define squ(x) ((x)*(x))

int k,n,idx;

struct Point

{

    int x[K];

    bool operator <(const Point& rhs) const{

        return x[idx] < rhs.x[idx];

    }

    void print() const {

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

            printf("%d ",x[i]);

        printf("%d\n",x[k-]);

    }

}P[maxn];

#define fi first

#define se second

typedef pair<double,Point> HeapNode;

priority_queue<HeapNode> q;

struct kdTree

{

    Point Node[maxn<<];

    int son[maxn<<];

    #define lch (rt<<1)

    #define rch (rt<<1|1)

    void build(int l,int r,int rt = ,int dep = )

    {

        if(l>r) return;

        son[rt] = r-l;

        int x = lch,y = rch;

        son[x] = son[y] = -;

        idx = dep%k;

        int mid = (l+r)>>;

        nth_element(P+l,P+mid,P+r+);

        Node[rt] = P[mid];

        build(l,mid-,x,dep+);

        build(mid+,r,y,dep+);

    }

    Point qp; int qm;

    void query(int rt = ,int dep = )

    {

        if(!~son[rt]) return;

        HeapNode u(,Node[rt]);

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

            u.fi += squ(u.se.x[i]-qp.x[i]);

        int dim = dep%k, x = lch, y = rch;

        bool flag = false;

        if(qp.x[dim]>=Node[rt].x[dim]) swap(x,y);

        if(~son[x]) query(x,dep+);

        if(q.size()<qm) q.push(u),flag = true;

        else {

            if(u.fi<q.top().fi) q.pop(),q.push(u);

            if(squ(qp.x[dim]-Node[rt].x[dim])<q.top().fi) flag = true;

        }

        if(flag&&~son[y]) query(y,dep+);

    }

}kd;

const int M = ;

Point ans[M];

int main()

{

   // freopen("in.txt","r",stdin);

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

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

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

            scanf("%d",&P[i].x[j]);

        kd.build(,n-);

        int t; scanf("%d",&t);

        while(t--){

            for(int i = ; i < k; i++) scanf("%d",&kd.qp.x[i]);

            scanf("%d",&kd.qm);

            kd.query();

            printf("the closest %d points are:\n",kd.qm);

            int top = ;

            while(q.size()){

                ans[++top] = q.top().se;

                q.pop();

            }

            while(top){

                ans[top].print();

                top--;

            }

        }

    }

    return ;

}