poj4091The Closest M Points(KD-tree)

kd-tree参考资料：

http://my.oschina.net/keyven/blog/221792

http://blog.csdn.net/lsjseu/article/details/12344443

#include <cstdio>
#include <algorithm>
#include <queue>
#include <cstring>
using namespace std;

const int MAXD = 5, MAXT = 5010;
int n, k, m, cur=1, curd;

class point{
public:
    int a[MAXD];
}pa[MAXT], tp;
inline int sqr(int a){return a*a; }
int dis(point &a, point &b){
    int res=0;
    for(int i=0; i<k; ++i)
        res += sqr(a.a[i]-b.a[i]);
    return res;
}

class cmp{
public:
    int operator()(const int &i, const int &j){
        return dis(tp,pa[i])<dis(tp,pa[j]);
    }
}less_than;

int comp(const point &i, const point &j){// 比较curd维的大小
    return i.a[curd]<j.a[curd];
}    
priority_queue<int,vector<int>,cmp > max_heap;

class node{
public:
    int l,r,dim,lc,rc,idx;
}kdt[MAXT*2];// 线段树，范围：[l,r]
void build_kdt(int lf, int rt, int d){
    int m = (lf+rt)>>1;
    curd = d;
    nth_element(pa+lf, pa+m, pa+rt+1, comp);
    node &root = kdt[cur++];
    root.l = lf, root.r = rt, root.dim = d,    root.idx = m;
    if(lf<m){
        root.lc = cur;
        build_kdt(lf, m-1, (d+1)%k);
    }
    if(m<rt){
        root.rc = cur;
        build_kdt(m+1, rt, (d+1)%k);
    }
}
void query(int r){
    if(r<=0) return;
    node &rt = kdt[r];
    // 访问当前结点
    if(max_heap.size()<m)
        max_heap.push(rt.idx);
    else if(less_than(rt.idx, max_heap.top())){
        max_heap.pop();
        max_heap.push(rt.idx);
    }
    int fir, sec;
    if(tp.a[rt.dim]<pa[rt.idx].a[rt.dim])
        fir = rt.lc, sec = rt.rc;
    else fir = rt.rc, sec = rt.lc;
    query(fir);// 访问目标所在的结点的子树，并进行是否在另一子树的判断
    if(max_heap.size()<m || dis(pa[max_heap.top()], tp)>=sqr(tp.a[rt.dim]-pa[rt.idx].a[rt.dim])){
        query(sec);
    }
}

void print(){// 从小到大输出最大堆
    if(!max_heap.empty()){
        int i = max_heap.top(); max_heap.pop();
        print();
        for(int j=0; j<k; ++j)
            printf("%d ", pa[i].a[j]);
        printf("
");
    }
}

int main(){
    // freopen("in.txt", "r", stdin);
    int t;
    while(scanf("%d%d", &n, &k) == 2){
        for(int i=0; i<n; ++i){
            for(int j=0; j<k; ++j)
                scanf("%d", &pa[i].a[j]);
        }
        memset(kdt, 0, sizeof(kdt));
        cur = 1;
        build_kdt(0,n-1,0);
        scanf("%d", &t);
        while(t--){
            for(int i=0; i<k; ++i) scanf("%d", &tp.a[i]);
            scanf("%d", &m);
            printf("the closest %d points are:
", m);
            while(!max_heap.empty()) max_heap.pop();
            query(1);
            print();
        }
    }
    return 0;
}