bzoj3053 The Closest M Points（多维KDtree）

题目链接

分析：
多维KDtree
用优先队列维护前K近的点
注意优先队列是“从大到小”进行排序的

tip

update的时候写错了max和min，查了好久。。。orz
多组数据！！！

#include<cstdio>
#include<cstring>
#include<iostream>
#include<algorithm>
#include<queue>

using namespace std;

const int INF=1e9;
const int N=500010;
struct node{
    int d[5],mx[5],mn[5],lc,rc,dis;
    bool operator <(const node &a) const    //从大到小 
    {
        return dis<a.dis;
    }
};
node t[N],ans[N];
int n,K,root,x[5],m;
priority_queue<node> Q;

void update(int bh)
{
    int lc=t[bh].lc;
    int rc=t[bh].rc;
    if (lc)
    {
        for (int i=0;i<K;i++)
        {
            t[bh].mn[i]=min(t[bh].mn[i],t[lc].mn[i]);
            t[bh].mx[i]=max(t[bh].mx[i],t[lc].mx[i]);
        }
    }
    if (rc)
    {
        for (int i=0;i<K;i++)
        {
            t[bh].mn[i]=min(t[bh].mn[i],t[rc].mn[i]);
            t[bh].mx[i]=max(t[bh].mx[i],t[rc].mx[i]);
        }
    }
}

int build(int l,int r,int D)
{
    int mid=(l+r)>>1;
    for (int i=l;i<=r;i++) t[i].dis=t[i].d[D];
    nth_element(t+l,t+mid+1,t+r+1);
    for (int i=0;i<K;i++)
        t[mid].mn[i]=t[mid].mx[i]=t[mid].d[i];
    if (l!=mid) t[mid].lc=build(l,mid-1,(D+1)%K);
    if (r!=mid) t[mid].rc=build(mid+1,r,(D+1)%K);
    update(mid);
    return mid;
}

int dis(int bh)
{
    int d=0;
    for (int i=0;i<K;i++)
    {
        if (x[i]<t[bh].mn[i]) d+=(t[bh].mn[i]-x[i])*(t[bh].mn[i]-x[i]);
        if (x[i]>t[bh].mx[i]) d+=(x[i]-t[bh].mx[i])*(x[i]-t[bh].mx[i]);
    }
    return d;
}

void ask(int bh)
{
    int d0=0,dl,dr;
    for (int i=0;i<K;i++) d0+=(t[bh].d[i]-x[i])*(t[bh].d[i]-x[i]); 
    t[bh].dis=d0;

    if (Q.size()<m) Q.push(t[bh]);   //少于m个点直接入队 
    else
    {
        if (Q.top().dis>d0)
        {
            Q.pop();
            Q.push(t[bh]);
        }
    }

    if (t[bh].lc) dl=dis(t[bh].lc);
    else dl=INF;
    if (t[bh].rc) dr=dis(t[bh].rc);
    else dr=INF;
    if (dl<dr)
    {
        if (dl<Q.top().dis||Q.size()<m) ask(t[bh].lc);
        if (dr<Q.top().dis||Q.size()<m) ask(t[bh].rc);
    }
    else
    {
        if (dr<Q.top().dis||Q.size()<m) ask(t[bh].rc);
        if (dl<Q.top().dis||Q.size()<m) ask(t[bh].lc);
    }
}

int main()
{
    while (scanf("%d%d",&n,&K)!=EOF)
    {
        memset(t,0,sizeof(t));
        for (int i=1;i<=n;i++)
            for (int j=0;j<K;j++) scanf("%d",&t[i].d[j]);
        root=build(1,n,0);
        int T;
        scanf("%d",&T);
        while (T--)
        {
            for (int i=0;i<K;i++) scanf("%d",&x[i]); 
            scanf("%d",&m);

            while (!Q.empty()) Q.pop();   //优先队列，从大到小(队首最大) 
            ask(root);
            for (int i=m;i>=1;i--) 
                ans[i]=Q.top(),Q.pop();
            printf("the closest %d points are:\n",m);
            for (int i=1;i<=m;i++)
            {
                for (int j=0;j<K-1;j++)
                    printf("%d ",ans[i].d[j]);
                printf("%d\n",ans[i].d[K-1]);
           } 
        }
    }
    return 0;
}