K - Finding Hotels HDU - 5992

kd树模板题目，建立kd树，对于每个输入进行搜索，在原来的基础上加入价格比对就可以了。

kd树其实就是一种暴力＋减枝的操作。

#include <algorithm>
#include <bitset>
#include <cassert>
#include <climits>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <deque>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <math.h>
#include <set>
#include <vector>
using namespace std;
#define Max_N (200000 + 100)
#define Max_M (20000 + 100)
#define ll long long int
const ll INF = 1e17;
struct P{
    int p[2], c, id;
};
P h[Max_N];
#define de 2
double var[de];
int split[Max_N];
int CM;
bool cmp(const P &a, const P &b)
{
    return a.p[CM] < b.p[CM];
}
void KD_build(int l, int r)
{
    if (l >= r) return;
    int mid = l + (r - l) / 2;
    for (int i = 0; i < de; i++) {
        double ave = 0;
        for (int j = l; j <= r; j++) {
            ave += h[j].p[i];
        }
        ave /= (r - l + 1);
        var[i] = 0;
        for (int j = l; j <= r; j++) {
            var[i] += (h[j].p[i] - ave) * (h[j].p[i] - ave);
        }
        var[i] /= (r - l + 1);
    }
    split[mid] = -1;
    double ms = -1;
    for (int i = 0; i < de; i++) {
        if (var[i] > ms) {
            split[mid] = i;
            ms = var[i];
        }
    }
    CM = split[mid];
    nth_element(h+l,h+mid,h+r+1,cmp);
    KD_build(l,mid-1);
    KD_build(mid+1,r);
    
}
int ans_i;
ll ansdis;
void KD_query(int l, int r, P s)
{
    int mid = l + (r - l) / 2;
    if (l > r) return;
    long long int dis = 0;
    for (int i = 0; i < de; i++) {
        dis += (ll)(s.p[i] - h[mid].p[i]) * (s.p[i] - h[mid].p[i]);
    }
    if (h[mid].c <= s.c) {
        if (ansdis == dis && h[mid].id < h[ans_i].id) {
            ans_i = mid;
        }
        if (ansdis > dis) {
            ansdis = dis;
            ans_i = mid;
        }
    }
    int d = split[mid];
    ll radius = (ll)(s.p[d] - h[mid].p[d]) * (s.p[d] - h[mid].p[d]);
    if (s.p[d] < h[mid].p[d]) {
        KD_query(l, mid-1, s);
        if (ansdis >= radius) {
            KD_query(mid+1, r, s);
        }
    }
    else {
        KD_query(mid+1, r, s);
        if (ansdis >= radius) {
            KD_query(l, mid-1, s);
        }
    }
}
int main()
{
    int T, n, m;
    scanf("%d", &T);
    while (T--) {
        scanf("%d%d", &n, &m);
        for (int i = 0; i < n; i++) {
            scanf("%d%d%d", &h[i].p[0], &h[i].p[1], &h[i].c);
            h[i].id = i;
        }
        KD_build(0, n-1);
        P s;
        for (int i = 0; i < m; i++) {
            scanf("%d%d%d", &s.p[0], &s.p[1], &s.c);
            ansdis = INF;
            ans_i = -1;
            KD_query(0, n-1, s);
            printf("%d %d %d\n", h[ans_i].p[0], h[ans_i].p[1], h[ans_i].c);
        }
    }
    return 0;
}