hdu 4347 The Closest M Points

KD tree 模板 ，多组样例不用初始化

#include <bits/stdc++.h>
using namespace std;
typedef int lint;
typedef long long LL;
const int maxn = 500005;
const int maxk = 10;
lint K,D;
struct Point{
    lint d[maxk],mn[maxk],mx[maxk],lc,rc;
    bool operator < ( const Point& b )const{
        return d[D] < b.d[D];
    }
}p[maxn],o;
LL Q( Point& a,Point& b ){
    LL res = 0;
    for( lint i = 0;i < K;i++ ){
        res += (LL)( (LL)a.d[i]-b.d[i] )*( a.d[i]-b.d[i] );
    }
    return res;
}
lint cnt;
priority_queue<pair<LL,Point>> que;
struct KD{
    Point t[maxn];
    lint rt;
    void push_up( lint x ){
        for( lint i = 0;i < K;i++ ) {
            if( t[x].lc ) {
                t[x].mx[i] = max(t[x].lc ? t[t[x].lc].mx[i] : -0x3f3f3f3f, t[x].mx[i] );
                t[x].mn[i] = min(t[x].lc ? t[t[x].lc].mn[i] : 0x3f3f3f3f,  t[x].mn[i] );
            }
            if( t[x].rc ){
                t[x].mx[i] = max(t[x].mx[i], t[x].rc ? t[t[x].rc].mx[i] : -0x3f3f3f3f);
                t[x].mn[i] = min(t[x].mn[i], t[x].rc ? t[t[x].rc].mn[i] : 0x3f3f3f3f);
            }
        }
    }
    lint build( lint k,lint l,lint r ){
        if( l > r ) return 0;
        lint mid = l+r>>1;
        D = k;
        nth_element( p+l, p+mid, p+r+1 );
        t[mid] = p[mid];
        for( lint i = 0;i < K;i++ )t[mid].mn[i] = t[mid].mx[i] = t[mid].d[i];
        t[mid].lc = build( (k+1)%K,l,mid-1 );
        t[mid].rc = build( (k+1)%K,mid+1,r );
        push_up(mid);
        return mid;
    }
    LL guess( lint x ){
        if(!x) return 0x3f3f3f3f;
        LL res = 0;
        for( lint i = 0;i < K;i++ ){
            if( o.d[i] < t[x].mn[i] ){
                res += (LL)((LL)o.d[i]-t[x].mn[i])*( o.d[i]-t[x].mn[i] );
            }
            else if( o.d[i] > t[x].mx[i] ){
                res += (LL)( (LL)o.d[i]-t[x].mx[i] )*( o.d[i]-t[x].mx[i] );
            }
        }
        return res;
    }
    void query( lint x ){
        if( que.size() < cnt ) que.push( make_pair(Q( o,t[x] ),t[x]) );
        else if( que.top().first > Q( o,t[x] ) ){
            que.pop();
            que.push( make_pair(Q( o,t[x] ),t[x]) );
        }
        LL ans1 = guess(t[x].lc);
        LL ans2 = guess( t[x].rc );
        if( ans1 < ans2 ){
            if( t[x].lc &&ans1 < que.top().first ){
                query( t[x].lc );
            }else if( t[x].lc&&que.size() < cnt ){
                query(t[x].lc);
            }
            if( t[x].rc&&ans2 < que.top().first ){
                query(t[x].rc);
            }else if( t[x].rc&&que.size() < cnt ){
                query(t[x].rc);
            }
        }else{
            if( t[x].rc&&ans2 < que.top().first ){
                query(t[x].rc);
            }else if( t[x].rc&&que.size() < cnt ){
                query(t[x].rc);
            }
            if( t[x].lc &&ans1 < que.top().first ){
                query( t[x].lc );
            }else if( t[x].lc&&que.size() < cnt ){
                query(t[x].lc);
            }
        }
    }
}T;
int main(){
    lint n;
    while (2 == scanf("%d%d",&n,&K)) {
        for (lint i = 1; i <= n; i++) {
            for (lint j = 0; j < K; j++) {
                scanf("%d", &p[i].d[j]);
            }
        }
        T.rt = T.build(0, 1, n);
        lint m;
        scanf("%d", &m);
        for (lint i = 1; i <= m; i++) {
            for (lint j = 0; j < K; j++) {
                scanf("%d", &o.d[j]);
            }
            scanf("%d", &cnt);
            printf("the closest %d points are:\n",cnt);
            T.query(T.rt);
            stack<Point> st;
            while (que.size()) {
                Point now = que.top().second;
                que.pop();
                st.push(now);
            }
            while(st.size()){
                Point now = st.top();
                st.pop();
                printf("%d", now.d[0]);
                for (lint k = 1; k < K; k++) {
                    printf(" %d", now.d[k]);
                }
                printf("\n");
            }
        }
    }
    return 0;
}