Kd树

从 Trac 迁移的文章

这是从旧校内 Wiki 迁移的文章,可能存在一些样式问题,您可以向 memset0 反馈。

原文章内容如下:

{{{
#define sqr(x) (x)*(x)  
int k,n,idx;   //k为维数,n为点数  
struct point  
{  
    int x[K];  
    bool operator < (const point &u) const  
    {  
        return x[idx]<u.x[idx];  
    }  
}po[N];  

typedef pair<double,point>tp;  
priority_queue<tp>nq;  

struct kdTree  
{  
    point pt[N<<2];  
    int son[N<<2];  

    void build(int l,int r,int rt=1,int dep=0)  
    {  
        if(l>r) return;  
        son[rt]=r-l;  
        son[rt*2]=son[rt*2+1]=-1;  
        idx=dep%k;  
        int mid=(l+r)/2;  
        nth_element(po+l,po+mid,po+r+1);  
        pt[rt]=po[mid];  
        build(l,mid-1,rt*2,dep+1);  
        build(mid+1,r,rt*2+1,dep+1);  
    }  
    void query(point p,int m,int rt=1,int dep=0)  
    {  
        if(son[rt]==-1) return;  
        tp nd(0,pt[rt]);  
        for(int i=0;i<k;i++) nd.first+=sqr(nd.second.x[i]-p.x[i]);  
        int dim=dep%k,x=rt*2,y=rt*2+1,fg=0;  
        if(p.x[dim]>=pt[rt].x[dim]) swap(x,y);  
        if(~son[x]) query(p,m,x,dep+1);  
        if(nq.size()<m) nq.push(nd),fg=1;  
        else  
        {
            if(nd.first<nq.top().first) nq.pop(),nq.push(nd);  
            if(sqr(p.x[dim]-pt[rt].x[dim])<nq.top().first) fg=1;  
        }
        if(~son[y]&&fg) query(p,m,y,dep+1); 

    }  
}kd;  
}}}
#define sqr(x) (x)*(x)  
int k,n,idx;   //k为维数,n为点数  
struct point  
{  
    int x[K];  
    bool operator < (const point &u) const  
    {  
        return x[idx]<u.x[idx];  
    }  
}po[N];  
typedef pair<double,point>tp;  
priority_queue<tp>nq;  
struct kdTree  
{  
    point pt[N<<2];  
    int son[N<<2];  
    void build(int l,int r,int rt=1,int dep=0)  
    {  
        if(l>r) return;  
        son[rt]=r-l;  
        son[rt*2]=son[rt*2+1]=-1;  
        idx=dep%k;  
        int mid=(l+r)/2;  
        nth_element(po+l,po+mid,po+r+1);  
        pt[rt]=po[mid];  
        build(l,mid-1,rt*2,dep+1);  
        build(mid+1,r,rt*2+1,dep+1);  
    }  
    void query(point p,int m,int rt=1,int dep=0)  
    {  
        if(son[rt]==-1) return;  
        tp nd(0,pt[rt]);  
        for(int i=0;i<k;i++) nd.first+=sqr(nd.second.x[i]-p.x[i]);  
        int dim=dep%k,x=rt*2,y=rt*2+1,fg=0;  
        if(p.x[dim]>=pt[rt].x[dim]) swap(x,y);  
        if(~son[x]) query(p,m,x,dep+1);  
        if(nq.size()<m) nq.push(nd),fg=1;  
        else  
        {
            if(nd.first<nq.top().first) nq.pop(),nq.push(nd);  
            if(sqr(p.x[dim]-pt[rt].x[dim])<nq.top().first) fg=1;  
        }
        if(~son[y]&&fg) query(p,m,y,dep+1); 
    }  
}kd;