HDU 5992/nowcoder 207K - Finding Hotels - [KDTree]

题目链接：http://acm.hdu.edu.cn/showproblem.php?pid=5992

题目链接：https://www.nowcoder.com/acm/contest/207/K

Problem Description
There are N hotels all over the world. Each hotel has a location and a price. M guests want to find a hotel with an acceptable price and a minimum distance from their locations. The distances are measured in Euclidean metric.

Input
The first line is the number of test cases. For each test case, the first line contains two integers N (N ≤ 200000) and M (M ≤ 20000). Each of the following N lines describes a hotel with 3 integers x (1 ≤ x ≤ N), y (1 ≤ y ≤ N) and c (1 ≤ c ≤ N), in which x and y are the coordinates of the hotel, c is its price. It is guaranteed that each of the N hotels has distinct x, distinct y, and distinct c. Then each of the following M lines describes the query of a guest with 3 integers x (1 ≤ x ≤ N), y (1 ≤ y ≤ N) and c (1 ≤ c ≤ N), in which x and y are the coordinates of the guest, c is the maximum acceptable price of the guest.

Output
For each guests query, output the hotel that the price is acceptable and is nearest to the guests location. If there are multiple hotels with acceptable prices and minimum distances, output the first one.

Sample Input
2
3 3
1 1 1
3 2 3
2 3 2
2 2 1
2 2 2
2 2 3
5 5
1 4 4
2 1 2
4 5 3
5 2 1
3 3 5
3 3 1
3 3 2
3 3 3
3 3 4
3 3 5

Sample Output
1 1 1
2 3 2
3 2 3
5 2 1
2 1 2
2 1 2
1 4 4
3 3 5

题意：

给出 $N$ 个酒店，每个酒店有坐标 $(x,y)$，价格为 $c$；同时又给出 $M$ 个人，同样每个人都有坐标 $(x,y)$，每个人能接收的最高酒店价格为 $c$；

每个人会选择不高于自己可接受价格的酒店中最近的那一个，如果同时有多个酒店同样最近，选择编号的那一个。

题解：

根据题目里的费用要求和同距离输出编号最小酒店两个要求，对KDTree模板进行一点小修改即可。

AC代码：

#include<bits/stdc++.h>
using namespace std;

const int maxn=2e5+;
const int maxdim=;

namespace KDTree
{
    int K;//维数
    inline double sqr(double x){return x*x;}
    struct Point
    {
        int x[maxdim];
        int c;
        int idx;
        double distance(const Point &oth)const
        {
            double ret=;
            for(int i=;i<K;i++) ret+=sqr(x[i]-oth.x[i]);
            return ret;
        }
        void input(int _idx)
        {
            idx=_idx;
            for(int i=;i<K;i++) scanf("%d",&x[i]);
            scanf("%d",&c);
        }
        void output()
        {
            for(int i=;i<K;i++) printf("%d ",x[i]);
            printf("%d\n",c);
        }
    };
    struct cmpx
    {
        int div;
        cmpx(const int &_div){div=_div;}
        bool operator()(const Point &a,const Point &b)
        {
            for(int i=;i<K;i++)
            {
                int k=(div+i)%K;
                if(a.x[k]!=b.x[k]) return a.x[k]<b.x[k];
            }
            return true;
        }
    };
    inline bool cmp(const Point &a,const Point &b,int div)
    {
        cmpx cp=cmpx(div);
        return cp(a,b);
    }

    struct Node //KDTree的节点
    {
        Point e;
        Node *lc,*rc;
        int div;
    }pool[maxn],*tail,*root;
    void init(){tail=pool;} //初始化KDTree
    Node* Build(Point *a,int l,int r,int div) //建树
    {
        if(l>=r) return NULL;
        Node *p=tail++;
        p->div=div;
        int mid=(l+r)/;
        nth_element(a+l,a+mid,a+r,cmpx(div));
        p->e=a[mid];
        p->lc=Build(a,l,mid,(div+)%K);
        p->rc=Build(a,mid+,r,(div+)%K);
        return p;
    }

    struct Qnode
    {
        Point p;
        double dist;
        Qnode(){}
        Qnode(Point _p,double _dist){p=_p; dist=_dist;}
        bool operator <(const Qnode &oth)const{return dist<oth.dist;}
    };
    priority_queue<Qnode> Q;
    void Search(const Point &p,Node *now,int div,int m) //在now节点子树下搜索p点的m近邻
    {
        if(now==NULL) return;
        if(cmp(p,now->e,div))
        {
            Search(p,now->lc,(div+)%K,m);
            if(Q.size()<m)
            {
                if(now->e.c <= p.c) Q.push(Qnode(now->e,p.distance(now->e)));
                Search(p,now->rc,(div+)%K,m);
            }
            else
            {
                if(p.distance(now->e) < Q.top().dist || (p.distance(now->e) == Q.top().dist && now->e.idx < Q.top().p.idx))
                {
                    if(now->e.c <= p.c)
                    {
                        Q.pop();
                        Q.push(Qnode(now->e,p.distance(now->e)));
                    }
                }
                if(sqr((now->e.x[div])-(p.x[div])) < Q.top().dist) Search(p,now->rc,(div+)%K,m);
            }
        }
        else
        {
            Search(p,now->rc,(div+)%K,m);
            if(Q.size()<m)
            {
                if(now->e.c <= p.c) Q.push(Qnode(now->e,p.distance(now->e)));
                Search(p,now->lc,(div+)%K,m);
            }
            else
            {
                if(p.distance(now->e) < Q.top().dist || (p.distance(now->e) == Q.top().dist && now->e.idx < Q.top().p.idx))
                {
                    if(now->e.c <= p.c)
                    {
                        Q.pop();
                        Q.push(Qnode(now->e,p.distance(now->e)));
                    }
                }
                if(sqr((now->e.x[div])-(p.x[div])) < Q.top().dist) Search(p,now->lc,(div+)%K,m);
            }
        }
    }

    void Search(const Point &p,int m) //搜索p点的m近邻
    {
        while(!Q.empty()) Q.pop();
        Search(p,root,,m);
    }
};

int n,m;
KDTree::Point p[maxn];

int main()
{
    KDTree::K=;
    int T;
    cin>>T;
    while(T--)
    {
        scanf("%d%d",&n,&m);
        for(int i=;i<n;i++) p[i].input(i);
        KDTree::init();
        KDTree::root=KDTree::Build(p,,n,);

        KDTree::Point o;
        for(int i=;i<=m;i++)
        {
            o.input(i);
            KDTree::Search(o,);
            o=KDTree::Q.top().p;
            o.output();
        }
    }
}

PS.由于只需要最近邻，所以其实并不需要那个优先队列，只不过懒得改了。