HDU 5992 Finding Hotels（KD树）题解

题意：n家旅店，每个旅店都有坐标x，y，每晚价钱z，m个客人，坐标x，y，钱c，问你每个客人最近且能住进去（非花最少钱）的旅店，一样近的选排名靠前的。

思路：KD树模板题

代码：

#include<set>
#include<map>
#include<stack>
#include<cmath>
#include<queue>
#include<vector>
#include<cstdio>
#include<cstring>
#include<iostream>
#include<algorithm>
typedef long long ll;
const int maxn = 200000 + 10;
const int seed = 131;
const ll MOD = 1e9 + 7;
const int INF = 0x3f3f3f3f;
using namespace std;

#define lson rt << 1
#define rson rt << 1 | 1
#define Pair pair<ll, Node>

int k, idx;   //维数k

struct Node
{
    int feature[3];     //定义属性数组
    int id;
    bool operator < (const Node &u) const
    {
        return feature[idx] < u.feature[idx];
    }
}_data[maxn], p;   //_data[]数组代表输入的数据
Node data[4 * maxn];    //data[]数组代表K-D树的所有节点数据
int flag[4 * maxn];      //用于标记某个节点是否存在，1表示存在，-1表示不存在

priority_queue<Pair> Q;     //队列Q用于存放离p最近的m个数据
ll Sqrt(ll a, ll b){    //欧几里得距离平方
    return (a - b) * 1LL * (a - b);
}

//建树步骤，参数dept代表树的深度
void Build(int l, int r, int rt, int dept)
{
    if(l > r) return;
    flag[rt] = 1;                   //表示编号为rt的节点存在
    flag[lson] = flag[rson] = -1;   //当前节点的孩子暂时标记不存在
    idx = dept % k;                 //按照编号为idx的属性进行划分
    int mid = (l + r) >> 1;
    nth_element(_data + l, _data + mid, _data + r + 1);   //nth_element()为STL中的函数 algorithm
    data[rt] = _data[mid];
    Build(l, mid - 1, lson, dept + 1);  //递归左子树
    Build(mid + 1, r, rson, dept + 1);  //递归右子树
}

//查询函数，寻找离p最近的m个特征属性
void Query(Node p, int m, int rt, int dept)
{
    if(flag[rt] == -1) return;   //不存在的节点不遍历
    Pair cur(0, data[rt]);       //获取当前节点的数据和到p的距离
    for(int i = 0; i < k; i++)   //欧几里得距离的平方
        cur.first += Sqrt((ll)data[rt].feature[i], (ll)p.feature[i]);
    int dim = dept % k;          //跟建树一样，这样能保证相同节点的dim值不变
    bool fg = 0;                 //用于标记是否需要遍历右子树
    int x = lson;
    int y = rson;
    if(p.feature[dim] >= data[rt].feature[dim]) //数据p的第dim个特征值大于等于当前的数据，则需要进入右子树
        swap(x, y);
    if(~flag[x]) Query(p, m, x, dept + 1);      //如果节点x存在，则进入子树继续遍历

    //以下是回溯过程，维护一个优先队列
    if(Q.size() < m)   //如果队列没有满，则继续放入
    {
        if(cur.second.feature[2] <= p.feature[2]) Q.push(cur);
        fg = 1;
    }
    else
    {
        if(cur.first < Q.top().first && cur.second.feature[2] <= p.feature[2])  //如果找到更小的距离，则用于替换队列Q中最大的距离的数据
        {
            Q.pop();
            Q.push(cur);
        }
        else if(cur.first == Q.top().first && cur.second.id < Q.top().second.id && cur.second.feature[2] <= p.feature[2]){
            Q.pop();
            Q.push(cur);
        }
        if(Sqrt((ll)p.feature[dim], (ll)data[rt].feature[dim]) < Q.top().first)
        {
            fg = 1;
        }
    }
    if(~flag[y] && fg)
        Query(p, m, y, dept + 1);
}
int main(){
    int T, n, m;
    k = 2;
    scanf("%d", &T);
    while(T--){
        scanf("%d%d", &n, &m);
        for(int i = 0; i < n; i++){
            for(int j = 0; j < 3; j++)
                scanf("%d", &_data[i].feature[j]);
            _data[i].id = i;
        }
        Build(0, n - 1, 1, 0);
        while(m--){
            while(!Q.empty()) Q.pop();
            for(int i = 0; i < 3; i++)
                scanf("%d", &p.feature[i]);
            Query(p, 1, 1, 0);
            p = Q.top().second;
            printf("%d %d %d
", p.feature[0], p.feature[1], p.feature[2]);
        }
    }
    return 0;
}

模板：

#define lson rt << 1
#define rson rt << 1 | 1
#define Pair pair<ll, Node>

int k, idx;   //维数k

struct Node
{
    int feature[3];     //定义属性数组
    int id;
    bool operator < (const Node &u) const
    {
        return feature[idx] < u.feature[idx];
    }
}_data[maxn], p;   //_data[]数组代表输入的数据
Node data[4 * maxn];    //data[]数组代表K-D树的所有节点数据
int flag[4 * maxn];      //用于标记某个节点是否存在，1表示存在，-1表示不存在

priority_queue<Pair> Q;     //队列Q用于存放离p最近的m个数据
ll Sqrt(ll a, ll b){    //欧几里得距离平方
    return (a - b) * 1LL * (a - b);
}

//建树步骤，参数dept代表树的深度
void Build(int l, int r, int rt, int dept)
{
    if(l > r) return;
    flag[rt] = 1;                   //表示编号为rt的节点存在
    flag[lson] = flag[rson] = -1;   //当前节点的孩子暂时标记不存在
    idx = dept % k;                 //按照编号为idx的属性进行划分
    int mid = (l + r) >> 1;
    nth_element(_data + l, _data + mid, _data + r + 1);   //nth_element()为STL中的函数 algorithm
    data[rt] = _data[mid];
    Build(l, mid - 1, lson, dept + 1);  //递归左子树
    Build(mid + 1, r, rson, dept + 1);  //递归右子树
}

//查询函数，寻找离p最近的m个特征属性
void Query(Node p, int m, int rt, int dept)
{
    if(flag[rt] == -1) return;   //不存在的节点不遍历
    Pair cur(0, data[rt]);       //获取当前节点的数据和到p的距离
    for(int i = 0; i < k; i++)   //欧几里得距离的平方
        cur.first += Sqrt((ll)data[rt].feature[i], (ll)p.feature[i]);
    int dim = dept % k;          //跟建树一样，这样能保证相同节点的dim值不变
    bool fg = 0;                 //用于标记是否需要遍历右子树
    int x = lson;
    int y = rson;
    if(p.feature[dim] >= data[rt].feature[dim]) //数据p的第dim个特征值大于等于当前的数据，则需要进入右子树
        swap(x, y);
    if(~flag[x]) Query(p, m, x, dept + 1);      //如果节点x存在，则进入子树继续遍历

    if(Q.size() < m)   //如果队列没有满，则继续放入
    {                  //注意，这里必须让fg=1，以后改时注意
        Q.push(cur);
        fg = 1;
    }
    else
    {
        if(cur.first < Q.top().first)  //如果找到更小的距离，则用于替换队列Q中最大的距离的数据
        {
            Q.pop();
            Q.push(cur);
        }
        if(Sqrt((ll)p.feature[dim], (ll)data[rt].feature[dim]) < Q.top().first)
        {
            fg = 1;
        }
    }
    if(~flag[y] && fg)
        Query(p, m, y, dept + 1);
}