hdu4347 kd-tree詢問k近點

親測直接寫暴力是過不了的，必須得加上神祕的剪枝技巧。

#include <iostream>
#include <cstdio>
#include <algorithm>
#include <cmath>
#include <cstring>
#include <queue>
#include <set>
#include <vector>
#include <map>
#define ll long long

using namespace std;

const int N=50007;
const int K=6;

int n,m;

struct point{
    int a[K];
    int div;  // 按哪個維度劃分
    bool lef;  // 是否是葉子節點
    ll dis;  // 離詢問點的距離。注意這個在讀入建樹時不會用到，在進入隊列時纔用到
    void print(){
        printf("%d",a[0]);
        for (int i=1;i<m;i++) 
            printf(" %d",a[i]);
        puts("");
    }
    bool operator < (const point &t) const{
        return dis<t.dis;
    }
    point(){}
    point(point &t,ll d){
        for (int i=0;i<m;i++) a[i]=t.a[i];
        dis=d;
    }
}p[N],tar;

int cmp_NO;
inline bool cmp(point x,point y){
    return x.a[cmp_NO]<y.a[cmp_NO];
}

inline ll dis(point x,point y){
    ll ret=0;
    for (int i=0;i<m;i++) 
        ret+=(x.a[i]-y.a[i])*(x.a[i]-y.a[i]);
    return ret;
}

inline void bulid_kdt(int L,int R,int d){
    if (L>R) return;
    int mid=(L+R)>>1;
    cmp_NO=d;
    nth_element(p+L,p+mid,p+R+1,cmp);
    p[mid].div=d;
    if (L==R){
        p[L].lef=true;
        return;
    }
    bulid_kdt(L,mid-1,(d+1)%m);
    bulid_kdt(mid+1,R,(d+1)%m);
}

priority_queue<point> que;
int num,nownum;
ll ans;

inline void find_kd(int L,int R){
    if (L>R) return;
    
    int mid=(L+R)>>1;
    ll d=dis(p[mid],tar);
    if (p[mid].lef){
        if (nownum<num){
            nownum++;
            que.push(point(p[mid],d));
            ans=max(ans,d);
        }
        else if (ans>d){
            que.pop();
            que.push(point(p[mid],d));
            ans=que.top().dis;
        }
        return;
    }

    int t=tar.a[p[mid].div]-p[mid].a[p[mid].div];
    if (t>0){
        find_kd(mid+1,R);
        if (nownum<num){
            nownum++;
            que.push(point(p[mid],d));
            ans=max(ans,d);
            find_kd(L,mid-1);
        }
        else {
            if (ans>d){
                que.pop();
                que.push(point(p[mid],d));
                ans=que.top().dis;
            }
            if (ans>t*t) 
                find_kd(L,mid-1);
        }
    }
    else {
        find_kd(L,mid-1);
        if (nownum<num){
            nownum++;
            que.push(point(p[mid],d));
            ans=max(ans,d);
            find_kd(mid+1,R);
        }
        else{
            if (ans>d){
                que.pop();
                que.push(point(p[mid],d));
                ans=que.top().dis;
            }
            if (ans>t*t) 
                find_kd(mid+1,R);
        }
    }
}

inline void put(){
    if (que.empty()) return;
    point pp=que.top();
    que.pop();
    put();
    pp.print();
}

int main(){
    while (~scanf("%d%d",&n,&m)){
        for (int i=0;i<n;i++){
            for (int j=0;j<m;j++) 
                scanf("%d",&p[i].a[j]);
            p[i].lef=false;
        }

        bulid_kdt(0,n-1,0);  // 這一步相當於將原數組重新排了個序，先訪問到的點放在中間

        int q;
        scanf("%d",&q);
        while (q--){
            for (int i=0;i<m;i++)
                scanf("%d",&tar.a[i]);
            while (!que.empty()) que.pop();
            scanf("%d",&num);
            nownum=0;
            ans=-1;
            find_kd(0,n-1);
            printf("the closest %d points are:\n",num);
            put();
        }
    }
    return 0;
}