KDTree 板子

从杨哥哪里偷的板子， 存一下。

#include<bits/stdc++.h>
using namespace std;
#define Fopen freopen("_in.txt","r",stdin); freopen("_out.txt","w",stdout);
#define LL long long
#define ULL unsigned LL
#define fi first
#define se second
#define pb push_back
#define lson l,m,rt<<1
#define rson m+1,r,rt<<1|1
#define lch(x) tr[x].son[0]
#define rch(x) tr[x].son[1]
#define max3(a,b,c) max(a,max(b,c))
#define min3(a,b,c) min(a,min(b,c))
typedef pair<int,int> pll;
const int inf = 0x3f3f3f3f;
const int _inf = 0xc0c0c0c0;
const LL INF = 0x3f3f3f3f3f3f3f3f;
const LL _INF = 0xc0c0c0c0c0c0c0c0;
const LL mod =  (int)1e9+7;
const int N = 5e4 + 100;
int n,k,idx;
struct Node{
    int f[5];
    bool operator <(const Node&rhs)const{
        return f[idx]<rhs.f[idx];
    }
}a[N];
priority_queue<pair<double,Node> >q;
struct kdtree{
    bool vis[N<<2];
    Node data[N<<2];
    void Build(int l, int r, int rt, int dep){
        if(l > r)return ;
        vis[rt] = 1;
        vis[rt<<1] = vis[rt<<1|1] = 0;
        idx = dep%k;
        int m = l+r >> 1;
        nth_element(a+l, a+m, a+r+1);
        data[rt] = a[m];
        Build(l, m-1, rt<<1, dep+1);
        Build(m+1, r, rt<<1|1, dep+1);
    }
    void Query(Node p, int m, int rt, int dep){
        if(!vis[rt]) return ;
        pair<double,Node>cur(0,data[rt]);
        for(int i = 0; i < k; ++i)
            cur.fi += (cur.se.f[i] - p.f[i]) * (cur.se.f[i] - p.f[i]);
        int dim = dep % k;
        bool flag = 0;
        int x = rt << 1,y = rt << 1 | 1;
        if(p.f[dim] >= data[rt].f[dim]) swap(x,y);
        if(vis[x]) Query(p, m, x, dep+1);
        if(q.size() < m) q.push(cur),flag = 1;
        else{
            if(cur.fi<q.top().fi){
                q.pop();
                q.push(cur);
            }
            if((p.f[dim]-data[rt].f[dim])*(p.f[dim]-data[rt].f[dim])<q.top().fi)
                flag=1;
        }
        if(vis[y] && flag) Query(p,m,y,dep+1);
    }
}kd;
int main(){
    while(scanf("%d%d",&n,&k)!=EOF){
        for(int i = 0;i < n; ++i)
            for(int j = 0; j < k; ++j)
                scanf("%d", &a[i].f[j]);
        kd.Build(0,n-1,1,0);
        int t;
        scanf("%d", &t);
        while(t--){
            Node p;
            for(int i = 0;i < k; ++i)
                scanf("%d", &p.f[i]);
            int m;
            scanf("%d",&m);
            while(!q.empty())
                q.pop();
            kd.Query(p,m,1,0);
            printf("the closest %d points are:
", m);
            Node ans[25];
            for(int i = 0; !q.empty(); ++i){
                ans[i] = q.top().se;
                q.pop();
            }
            for(int i = m - 1; i >= 0; --i)
                for(int j = 0; j < k; ++j)
                    printf("%d%c", ans[i].f[j], " 
"[j==k-1]);
        }
    }
    return 0;
}