【K-D树 在限制条件下求最近欧几里德距离】HDU - 5992 Finding Hotels

Problem Description

给你n个酒店，m个人。
分别给你每个酒店的坐标，和酒店的价格。
分别给你每个人的坐标，和人最多能够承受的价格。
对于每个人，让你求这个人能够承受价格的酒店，那个距离他最近。

思路：

求最近距离的时候，加个条件限制即可。

#include<bits/stdc++.h>
using namespace std;
#define ll long long
const int DIM = 3;//三维，第三维是价格
const int MAX = 200055;
const ll inf = 0x3f3f3f3f3f3f3f3f;
struct node
{
    int l, r;
    int id;//用来记录输入顺序
    int d[DIM], maxn[DIM], minn[DIM];
    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;
}
inline ll sqr(int x)
{
    return (ll)x * (ll)x;
}
inline ll dis(int o, int k)//求两点间距离的平方
{
    return sqr(tree[o].d[0] - tree[k].d[0]) + sqr(tree[o].d[1] - tree[k].d[1]);
}
inline ll partionMin(int o, int k)//k->o点 可能的最近距离
{
    if(tree[o].minn[2] > tree[k].d[2]) return inf;//价格不满足，距离返回最远
    ll red = 0;
    for(int i = 0; i < DIM - 1; i++)
    {
        if(tree[k].d[i] > tree[o].maxn[i]) red += sqr(tree[o].maxn[i] - tree[k].d[i]);
        if(tree[k].d[i] < tree[o].minn[i]) red += sqr(tree[o].minn[i] - tree[k].d[i]);
    }
    return red;
}
ll ans;
int ans_id, ans_node;
void query(int o, int k)//求满足条件最近距离的酒店
{
   // printf("%d %d %d\n", tree[o].d[0], tree[o].d[1], tree[o].d[2]);
    ll dm = dis(o, k);//o k两点距离的平方
    if(tree[o].d[2] <= tree[k].d[2] && (dm < ans || (dm == ans && tree[o].id < ans_id)))//满足条件，更新
    {
        ans_id = tree[o].id;
        ans = dm;
        ans_node = 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 <= ans) query(tree[o].l, k);
        if(dr <= ans) query(tree[o].r, k);
    }
    else
    {
        if(dr <= ans) query(tree[o].r, k);
        if(dl <= ans) query(tree[o].l, k);
    }
}
int main()
{
    int T, n, m, i, j;
    scanf("%d", &T);
    while(T--)
    {
        scanf("%d %d", &n, &m);
        for(i = 1; i <= n; i++){
            for(j = 0; j < DIM; j++)
            scanf("%d", &tree[i].d[j]);
            tree[i].id = i;
        }
        int root = build(1, n, 0);
        while(m--)
        {
            for(j = 0; j < DIM; j++)
                scanf("%d", &tree[n+1].d[j]);
            ans = inf;
            query(root, n+1);
          //  printf("%d\n", ansid);
            for(j = 0; j < DIM; j++)
            {
                if(j) printf(" ");
                printf("%d", tree[ans_node].d[j]);
            }
            printf("\n");
        }
    }
    return 0;
}