HDU The Closest M Points（KD-tree模板题）

题目链接

HDU 4347 The Closest M Points

分析

kd-tree模板题

AC code

#include<bits/stdc++.h>#define pb push_back#define mp make_pair#define PI acos(-1)#define fi first#define se second#define INF 0x3f3f3f3f#define INF64 0x3f3f3f3f3f3f3f3f#define random(a,b) ((a)+rand()%((b)-(a)+1))#define ms(x,v) memset((x),(v),sizeof(x))#define scint(x) scanf("%d",&x );#define scf(x) scanf("%lf",&x );#define eps 1e-10#define dcmp(x) (fabs(x) < eps? 0:((x) <0?-1:1))#define SQ(x) (x)*(x)using namespace std;typedef long long LL;typedef long double DB;typedef pair<int,int> Pair;const int maxn = 5e4+10;const int MOD = 1e9+7;int _idx;//比较维度struct KDNode{    const static int max_dims = 5;    int featrue[max_dims];    int size;//子树节点个数    int region[max_dims][2];//每个维度最大值最小值    int dim;    bool operator < (const KDNode& o)const{        return featrue[_idx]<o.featrue[_idx];    }};struct KDTree{    int dims;    KDNode Node[maxn];    KDNode data[maxn<<2];    bool flag[maxn<<2];    priority_queue<pair<int,KDNode> > Q;//查询结果队列    void build(int l,int r,int o,int dep,bool clc_region = false){        //最后一个参数表明是否记录区域大小        if(l>r)return;        _idx = dep % dims;        int lc = o<<1,rc = o<<1|1;        flag[o] = true;        flag[lc]=flag[rc] = 0;        int mid = (l+r) >> 1;        nth_element(Node+l,Node+mid,Node+r+1);        data[o] = Node[mid];data[o].dim = _idx;        // std::cout <<"node "<< o << '\n';        // std::cout << _idx << '\n';        // for(int i=0 ; i<dims ; ++i)std::cout << data[o].featrue[i] << ' ';std::cout  << '\n';        data[o].size = r-l+1;        if(clc_region){            for(int i=0 ; i<dims ; ++i){                _idx = i;                data[o].region[i][0] = min_element(Node+l,Node+r+1)->featrue[i];                data[o].region[i][1] = max_element(Node+l,Node+r+1)->featrue[i];            }            _idx = dep%dims;        }        build(l,mid-1,lc,dep+1,clc_region);        build(mid+1,r,rc,dep+1,clc_region);    }    void k_close(const KDNode& p,int k,int o){        if(!flag[o])return;        int dim = data[o].dim;        int lc = o<<1;int rc = o<<1|1;        if(p.featrue[dim] >data[o].featrue[dim])swap(lc,rc);        if(flag[lc])k_close(p,k,lc);        pair<int,KDNode> cur(0,data[o]);        for(int i=0 ; i<dims ; ++i)cur.fi+=SQ(p.featrue[i]-data[o].featrue[i]);        bool fg = false;//右子树遍历标志        if(Q.size() < k){            Q.push(cur);fg =1;        }else{            if(cur.fi < Q.top().fi){                Q.pop();Q.push(cur);            }            fg = SQ(p.featrue[dim]-data[o].featrue[dim]) < Q.top().fi;        }        if(flag[rc] && fg)k_close(p,k,rc);    }    int  check(int region[][2],int o){        //1表示相交        //－１表示全属于        //0表示不相交        if(!flag[o])return 0;        bool fg = true;        for(int i=0 ; i<dims ; ++i){            if(data[o].region[i][0] < region[i][0] || data[o].region[i][1] > region[i][1]){                fg = false;break;            }        }        int d = data[o].dim;        return fg?-1 : data[o].region[d][1] > region[d][0] || data[o].region[d][0]<region[d][1];    }    int find_size(int region[][2],int o){        //查找范围内的点数        //默认建树时有ｒｅｇｉｏｎ记录        if(!flag[o])return 0;        int ret =0;        bool fg =1 ;//当前点是否在范围内        for(int i=0 ; i<dims ; ++i)            if(data[o].featrue[i]<region[i][0]||data[o].featrue[i]>region[i][1]){                fg = 0;break;            }        ret += fg;        int lc = o<<1,rc = o<<1|1;        int lstate = check(region,lc),rstate = check(region,rc);        if(lstate ==-1)ret += data[lc].size;        else if(lstate == 1)ret += find_size(region,lc);        if(rstate ==-1)ret += data[rc].size;        else if(rstate == 1)ret += find_size(region,rc);        return ret;    }};KDTree kd;int main(){    // ios_base::sync_with_stdio(0);    // cin.tie(0);    // cout.tie(0);    int n,k;    while (scanf("%d%d",&n,&k )!=EOF) {        kd.dims = k;        for(int i=0 ; i<n ; ++i){            for(int j =0 ; j<k ; ++j)scanf("%d",&kd.Node[i].featrue[j]);        }        kd.build(0,n-1,1,0);        int m;        scanf("%d",&m );        while (m--) {            KDNode p;            for(int i=0 ; i<k ; ++i)scanf("%d",&p.featrue[i] );            int num =0;            scanf("%d",&num );            kd.k_close(p,num,1);            std::vector<KDNode> v;            while (!kd.Q.empty()) {                v.pb(kd.Q.top().se);kd.Q.pop();            }            printf("the closest %d points are:\n",num);            for(int i = v.size()-1 ; i>=0 ; --i){                for(int j=0 ; j<k ; ++j){                    printf("%d%c",v[i].featrue[j],j==k-1?'\n':' ' );                }            }        }    }    //std::cout << "time "<< clock()/1000 <<"ms"<< '\n';    return 0;}