【K-D树 K维最近距离的t个点】HDU - 4347 The Closest M Points

Problem Description

给你n个点，告诉你k维。
接下来有m个询问，每个询问，给你一个坐标，和t。
让你求距离这个坐标最近的t个点。

思路：

求最近距离的基础上，加个优先队列存t个点。

#include<bits/stdc++.h>
using namespace std;
#define ll long long
const int MAX = 50055;
const ll inf = 0x3f3f3f3f3f3f3f3f;
int DIM;
struct node
{
    int l, r;
    int d[5], maxn[5], minn[5];//最大是5维
    inline void maintain()//初始化
    {
        l = r = 0;
        for(int i = 0; i < DIM; i++)
            maxn[i] = minn[i] = d[i];
    }
}tree[MAX];
int D;
bool operator < (const node &x, const node &y)//重载从小到大
{
    return x.d[D] < y.d[D];
}
inline void Merge(int o)//归并，向上更新
{
    int son[2] = {tree[o].l, tree[o].r};
    for(int i = 0; i < 2; i++)
    {
        if(!son[i]) continue;
        for(int j = 0; j < DIM; j++)
        {
            tree[o].maxn[j] = max(tree[o].maxn[j], tree[son[i]].maxn[j]);
            tree[o].minn[j] = min(tree[o].minn[j], tree[son[i]].minn[j]);
        }
    }
}
int build(int l, int r, int now)//建树
{
    int mid = (l+r) >> 1;
    D = now;
    nth_element(tree+l, tree+mid, tree+r+1);
    tree[mid].maintain();
    if(l < mid) tree[mid].l = build(l, mid-1, (now+1)%DIM);
    if(r > mid) tree[mid].r = build(mid+1, r, (now+1)%DIM);
    Merge(mid);
    return mid;
}
ll sqr(int x)
{
    return (ll)x * (ll)x;
}
ll dis(int o, int k)//求两点间距离的平方
{
    ll sum = 0;
    for(int i = 0; i < DIM; i++)
    {
        sum += sqr(tree[o].d[i] - tree[k].d[i]);
    }
    return sum;
}
ll partionMin(int o, int k)//k->o这个子树上所有点，可能的最小值
{
    ll red = 0;
    for(int i = 0; i < DIM; i++)
    {
        if(tree[k].d[i] > tree[o].maxn[i]) red += sqr(tree[k].d[i] - tree[o].maxn[i]);
        if(tree[k].d[i] < tree[o].minn[i]) red += sqr(tree[o].minn[i] - tree[k].d[i]);
    }
    return red;
}
struct Queue//优先队列节点
{
    ll dist;
    int id;
    bool operator < (const Queue &b) const{//优先出来最大的
        return dist < b.dist;
    }
};
priority_queue<Queue> q;
void query(int o, int k)
{
    ll dm = dis(o, k);
   // printf("dm = %lld\n", dm);
    if(dm < q.top().dist) {//优先队列最大的 大于 dm，更新
        q.pop();
        q.push((Queue){dm, o});
    }
    ll dl = tree[o].l ? partionMin(tree[o].l, k) : inf;
    ll dr = tree[o].r ? partionMin(tree[o].r, k) : inf;
    if(dl < dr)//优化时间
    {
        if(dl < q.top().dist) query(tree[o].l, k);
        if(dr < q.top().dist) query(tree[o].r, k);
    }
    else
    {
        if(dr < q.top().dist) query(tree[o].r, k);
        if(dl < q.top().dist) query(tree[o].l, k);
    }
}
void input()//输出
{
    int id;
    if(!q.empty())
    {
        id = q.top().id;
        q.pop();
        input();
      //  printf("id = %d\n", id);
        for(int i = 0; i < DIM; i++)
        {
            if(i) printf(" ");
            printf("%d", tree[id].d[i]);
        }
        printf("\n");
    }
}
int main()
{
    int n, i, j, m, t;
    while(~scanf("%d %d", &n, &DIM))
    {
        for(i = 1; i <= n; i++)
            for(j = 0; j < DIM; j++)
                scanf("%d", &tree[i].d[j]);
        int root = build(1, n, 0);
        scanf("%d", &m);
        while(m--)
        {
            for(j = 0; j < DIM; j++)
                scanf("%d", &tree[n+1].d[j]);
            scanf("%d", &t);
            for(j = 0; j < t; j++)//初始化优先队列
            q.push((Queue){inf, 0});
            query(root, n+1);
            printf("the closest %d points are:\n", t);
            input();
        }
    }
    return 0;
}