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+;
const int N = 5e4 + ;
int n,k,idx;
struct Node{
    int f[];
    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<<];
    Node data[N<<];
    void Build(int l, int r, int rt, int dep){
        if(l > r)return ;
        vis[rt] = ;
        vis[rt<<] = vis[rt<<|] = ;
        idx = dep%k;
        int m = l+r >> ;
        nth_element(a+l, a+m, a+r+);
        data[rt] = a[m];
        Build(l, m-, rt<<, dep+);
        Build(m+, r, rt<<|, dep+);
    }
    void Query(Node p, int m, int rt, int dep){
        if(!vis[rt]) return ;
        pair<double,Node>cur(,data[rt]);
        for(int i = ; 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 = ;
        int x = rt << ,y = rt <<  | ;
        if(p.f[dim] >= data[rt].f[dim]) swap(x,y);
        if(vis[x]) Query(p, m, x, dep+);
        if(q.size() < m) q.push(cur),flag = ;
        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=;
        }
        if(vis[y] && flag) Query(p,m,y,dep+);
    }
}kd;
int main(){
    while(scanf("%d%d",&n,&k)!=EOF){
        for(int i = ;i < n; ++i)
            for(int j = ; j < k; ++j)
                scanf("%d", &a[i].f[j]);
        kd.Build(,n-,,);
        int t;
        scanf("%d", &t);
        while(t--){
            Node p;
            for(int i = ;i < k; ++i)
                scanf("%d", &p.f[i]);
            int m;
            scanf("%d",&m);
            while(!q.empty())
                q.pop();
            kd.Query(p,m,,);
            printf("the closest %d points are:\n", m);
            Node ans[];
            for(int i = ; !q.empty(); ++i){
                ans[i] = q.top().se;
                q.pop();
            }
            for(int i = m - ; i >= ; --i)
                for(int j = ; j < k; ++j)
                    printf("%d%c", ans[i].f[j], " \n"[j==k-]);
        }
    }
    return ;
}