1. 首页
2. KDTree
3. 本页

HDU 4347 The Closest M Points (KD树)

题意：求K维空间中给定一个点最邻近的M个点。

思路：KD树模板题，

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
typedef unsigned long long ULL;
typedef pair<int, int> PI;
typedef pair< PI, int> PII;
#define pb push_back
#define Key_value ch[ch[root][1]][0]
#define lson l,mid,rt<<1
#define rson mid+1,r,rt<<1|1
#define MS(x,y) memset(x,y,sizeof(x))
const double eps=1e-6;
const double pi=acos(-1.0);
const int mod=1e9+7;
const int INF=0x3f3f3f3f;
const int M = 1000100;
const int MAXN =50010;
const int DIM=10;
inline double sqr(double x)
{
    return x*x;
}
namespace KDTree
{
int k;
struct Point
{
    int x[DIM];
    double distance(const Point &b)const
    {
        double ret=0;
        for(int i=0; i<k; ++i)
            ret+=sqr(x[i]-b.x[i]);
        return ret;
    }
};
struct qnode
{
    Point p;
    double dis;
    qnode(Point _p,double _dis)
    {
        p=_p;
        dis=_dis;
    }
    bool operator <(const qnode &b)const
    {
        return dis<b.dis;
    }
};
priority_queue<qnode>q;
struct cmpx
{
    int div;
    cmpx(const int &_div)
    {
        div=_div;
    }
    bool operator ()(const Point &a,const Point &b)
    {
        for(int i=0; i<k; ++i)
            if(a.x[(div+i)%k]!=b.x[(div+i)%k])
                return a.x[(div+i)%k]<b.x[(div+i)%k];
        return true;
    }
};
bool cmp(const Point &a,const Point &b,int div)
{
    cmpx cp=cmpx(div);
    return cp(a,b);
}
struct Node
{
    Point e;
    Node *lc,*rc;
    int div;
} pool[MAXN],*tail,*root;
void init()
{
    tail=pool;
}
Node* build(Point *a,int l,int r,int div)
{
    if(l>=r) return NULL;
    Node *p=tail++;
    p->div=div;
    int mid=(l+r)/2;
    nth_element(a+l,a+mid,a+r,cmpx(div));
    p->e=a[mid];
    p->lc=build(a,l,mid,(div+1)%k);
    p->rc=build(a,mid+1,r,(div+1)%k);
    return p;
}
void search(Point p,Node *x,int div,int m)
{
    if(!x) return;
    if(cmp(p,x->e,div))
    {
        search(p,x->lc,(div+1)%k,m);
        if(q.size()<m)
        {
            q.push(qnode(x->e,p.distance(x->e)));
            search(p,x->rc,(div+1)%k,m);
        }
        else
        {
            if(p.distance(x->e) < q.top().dis)
            {
                q.pop();
                q.push(qnode(x->e,p.distance(x->e)));
            }
            if(sqr(x->e.x[div]-p.x[div])<q.top().dis)
                search(p,x->rc,(div+1)%k,m);
        }
    }
    else
    {
        search(p,x->rc,(div+1)%k,m);
        if(q.size()<m)
        {
            q.push(qnode(x->e,p.distance(x->e)));
            search(p,x->lc,(div+1)%k,m);
        }
        else
        {
            if(p.distance(x->e) < q.top().dis)
            {
                q.pop();
                q.push(qnode(x->e,p.distance(x->e)));
            }
            if(sqr(x->e.x[div]-p.x[div])<q.top().dis)
                search(p,x->lc,(div+1)%k,m);
        }
    }
}
void search(Point p,int m)
{
    while(!q.empty())q.pop();
    search(p,root,0,m);
}
};

KDTree::Point p[MAXN];

int main()
{
    int n,k;
    while(~scanf("%d%d",&n,&k))
    {
        KDTree::k=k;
        for(int i=0; i<n; ++i)
        {
            for(int j=0; j<k; ++j)
            {
                scanf("%d",&p[i].x[j]);
            }
        }
        KDTree::init();
        KDTree::root=KDTree::build(p,0,n,0);
        KDTree::Point tmp;
        int Q;
        scanf("%d",&Q);
        while(Q--)
        {
            for(int i=0; i<k; ++i)
                scanf("%d",&tmp.x[i]);
            int m;
            scanf("%d",&m);
            KDTree::search(tmp,m);
            int cnt=0;
            while(!KDTree::q.empty())
            {
                p[cnt++]=KDTree::q.top().p;
                KDTree::q.pop();
            }
            printf("the closest %d points are:\n",m);
            for(int i=0; i<cnt; ++i)
            {
                for(int j=0; j<k; ++j)
                    printf("%d%c",p[cnt-1-i].x[j],j==k-1?'\n':' ');
            }
        }
    }
    return 0;
}