BZOJ 3053 The Closest M Points

【题目分析】

典型的KD-Tree例题，求k维空间中的最近点对，只需要在判断的过程中加上一个优先队列，就可以了。

【代码】

#include <cstdio>
#include <cstring>
#include <cstdlib>
#include <cmath>

#include <set>
#include <map>
#include <string>
#include <algorithm>
#include <vector>
#include <iostream>
#include <queue>

using namespace std;

#define maxn 1000005
#define inf (0x3f3f3f3f)
#define mk(a,b) make_pair(a,b)

int read()
{
    int x=0,f=1; char ch=getchar();
    while (ch<'0'||ch>'9') {if (ch=='-') f=-1; ch=getchar();}
    while (ch>='0'&&ch<='9') {x=x*10+ch-'0'; ch=getchar();}
    return x*f;
}

int D,rt,n,base,m,x;

struct node{
    int d[6],mx[6],mn[6],l,r;
    int operator [] (int x) {return d[x];}
    void init() {for (int i=0;i<base;++i) d[i]=read();}
}t[maxn],now;

bool operator < (node a,node b) {return a[D]<b[D];}

void update(int k)
{
    for (int i=0;i<base;++i)
    {
        t[k].mn[i]=min(t[k][i],min(t[t[k].l].mn[i],t[t[k].r].mn[i]));
        t[k].mx[i]=max(t[k][i],max(t[t[k].l].mx[i],t[t[k].r].mx[i]));
    }
}

int build(int l,int r,int dir)
{
    D=dir;
    int mid=(l+r)/2;
    nth_element(t+l,t+mid,t+r+1);
    for (int i=0;i<base;++i) t[mid].mn[i]=t[mid].mx[i]=t[mid][i];
    if (l<mid) t[mid].l=build(l,mid-1,(dir+1)%base); else t[mid].l=0;
    if (r>mid) t[mid].r=build(mid+1,r,(dir+1)%base); else t[mid].r=0;
    update(mid);
    return mid;
}

int dis(node a,node b)
{
    int ret=0;
    for (int i=0;i<base;++i)
        ret+=(a[i]-b[i])*(a[i]-b[i]);
    return ret;
}

pair <int,int> pa;
int ans[maxn];
priority_queue <pair<int,int> > q;

int getdis(int k)
{
    if (!k) return inf;
    int ret=0;
    for (int i=0;i<base;++i)
        if (now[i]<t[k].mn[i]) ret+=(t[k].mn[i]-now[i])*(t[k].mn[i]-now[i]);
    for (int i=0;i<base;++i)
        if (now[i]>t[k].mx[i]) ret+=(now[i]-t[k].mx[i])*(now[i]-t[k].mx[i]);
    return ret;
}

void query(int k)
{
    if (!k) return ;
    int dl=getdis(t[k].l),dr=getdis(t[k].r),d0=dis(t[k],now);
    if (d0<q.top().first) {q.pop(); q.push(mk(d0,k));}
    if (dl<dr)
    {
        if (dl<q.top().first) query(t[k].l);
        if (dr<q.top().first) query(t[k].r);
    }
    else
    {
        if (dr<q.top().first) query(t[k].r);
        if (dl<q.top().first) query(t[k].l);
    }
}

int main()
{
    while (scanf("%d%d",&n,&base)!=EOF)
    {
//      memset(t,0,sizeof t);
//      n=read(); base=read();
        for (int i=0;i<base;++i) t[0].mn[i]=inf,t[0].mx[i]=-inf;
        for (int i=1;i<=n;++i) t[i].init();
        rt=build(1,n,0);
        m=read();
        while (m--)
        {
            now.init(); x=read();
            for (int i=0;i<x;++i) q.push(mk(inf,0));
            query(rt);
            printf("the closest %d points are:\n",x);
            for (int i=x;i>=1;--i) ans[i]=q.top().second,q.pop();
            for (int i=1;i<=x;++i)
                for (int j=0;j<base;++j)
                    printf("%d%c",t[ans[i]][j],j==base-1?'\n':' ');
        }
    }
}