HDU 2966 In case of failure (k-d树)

题意：给出n个点，找出每个点距离他们最近的点的距离的平方。

思路：k-d树的模板题，先将所有点建立好k-d树，然后在依次求一下距离他们最近的点，如果找到的点是他们自己的时候设为距离为无穷大即可。

///k-d树
#include<cstdio>
#include<cstring>
#include<cmath>
#include<algorithm>
typedef long long ll;
const int maxn = 2 * 1e5 + 10;
const ll INF = 3 * 1e18;
const double eps = 1e-6;
using namespace std;

int k, n, nodecnt, t; ///点的维数,个数
struct point {
    ll data[5];
    int split; ///分裂维
    void input() { for(int i = 0; i < k; i++) scanf("%lld", &data[i]); }
    bool operator < (point p) const {
        int i = split;
        if(data[i] != p.data[i]) return data[i] < p.data[i];
        for(int j = 0; j < k; j++) {
            if(data[j] == p.data[j]) continue;
            return data[j] < p.data[j];
        }
        return false;
    }
} p[maxn], ot[maxn];
struct node {
    int son[2];
    point p;
    int split; ///分裂维
} kd[maxn];

///建立k-d树
void build(int o, int l, int r) {
    if(r < l) return ;
    int nn = r - l + 1, sp;
    double EX[5], EXX[5];
    memset(EX, 0, sizeof EX);     ///E(X)
    memset(EXX, 0, sizeof EXX);   ///E(X^2)
    for(int i = l; i <= r; i++) {
        for(int j = 0; j < k; j++) {
            EX[j] += (double)p[i].data[j] / nn;
            EXX[j] += (double)p[i].data[j] * p[i].data[j] / nn;
        }
    }
    ///D(X) = E(X^2) - E(X)^2，将方差最大的维作为分裂维
    double dx = -1;
    for(int i = 0; i < k; i++) {
        double d = EXX[i] - EX[i] * EX[i];
        if(dx < d) { dx = d; sp = i; }
    }
    kd[o].split = sp;
    for(int i = l; i <= r; i++) {
        p[i].split = sp;
    }
    int mid = (l + r) >> 1;
    ///sort(p + l, p + r + 1);
    nth_element (p + l, p + mid, p + r + 1);  ///O(n)
    for(int i = 0; i < k; i++) {
        kd[o].p.data[i] = p[mid].data[i];
    }
    if(mid - 1 >= l) {
        nodecnt++; kd[o].son[0] = nodecnt;
        build(kd[o].son[0], l, mid - 1);
    }
    if(mid + 1 <= r) {
        nodecnt++; kd[o].son[1] = nodecnt;
        build(kd[o].son[1], mid + 1, r);
    }
}
///两点的距离平方
ll dis(point a, point b) {
    ll d = 0;
    for(int i = 0; i < k; i++) {
        d += (a.data[i] - b.data[i]) * (a.data[i] - b.data[i]);
    }
    return d;
}

///寻找距离点p最近的点
ll mindistance(int o, point po) {
    if(!o) return INF;
    int s = 0;
    ///寻找进入的子树
    if(kd[o].p.data[kd[o].split] != po.data[kd[o].split]) {
        s = kd[o].p.data[kd[o].split] > po.data[kd[o].split] ? 0 : 1;
    } else {
        for(int i = 0; i < k; i++) {
            if(kd[o].p.data[i] == po.data[i]) continue;
            s = kd[o].p.data[i] > po.data[i] ? 0 : 1;
            break;
        }
    }
    ///查询点的位置
    int nxt = kd[o].son[s];
    int other = kd[o].son[s ^ 1];
    ll di = dis(kd[o].p, po), dr = 0;
    if(!di) di = INF;
    ll dl = mindistance(nxt, po);
    if(dl) di = min(di, dl);
    if(!di) di = INF;
    double g = sqrt(di);
    
    ///是否要进入另一棵子树
    ll f = s ? -1 : 1;
    ll lo = kd[o].p.data[kd[o].split]; ///当前节点分裂维的坐标
    double lp = po.data[kd[o].split] + f * sqrt(di);
    if(s == 0 && lp + eps > lo) dr = mindistance(other, po);
    if(s == 1 && lp - eps < lo) dr = mindistance(other, po);
    if(di && dr) di = min(di, dr);
    else if(!di) di = dr;
    if(!di) di = INF;
    return di;
}

int main() {
    k = 2;
    scanf("%d", &t);
    while(t--) {
        scanf("%d", &n);  nodecnt = 1;
        for(int i = 0; i < maxn; i++) {
            kd[i].son[0] = kd[i].son[1] = 0;
        }
        for(int i = 0; i < n; i++) {
            p[i].input();
            ot[i].data[0] = p[i].data[0];
            ot[i].data[1] = p[i].data[1];
        }
        build(1, 0, n - 1);
        for(int i = 0; i < n; i++) {
            printf("%lld\n", mindistance(1, ot[i]));
        }
    }
    return 0;
}