BZOJ3053 The Closest M Points-程序员宅基地

裸的KD-tree，还是在估计要不要进入子树的时候判断一下就好了，剩下都一样

判断的方法就是看现在答案个数是否小于k，和答案是否会经过分割线。

 

/**************************************************************
    Problem: 3053
    User: rausen
    Language: C++
    Result: Accepted
    Time:2164 ms
    Memory:3572 kb
****************************************************************/
 
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <queue>
 
using namespace std;
typedef long long ll;
 
const int Cnt_kd = 1e5 + 5;
const int Di = 5;
 
struct kd_node {
    kd_node *s[2];
    int p[Di];
} *kd_root, mempool[Cnt_kd], *cnt_kd, *null, *ans[25];
 
struct data {
    kd_node *from;
    ll dis;
    data() {}
    data(kd_node *_f, ll _d) : from(_f), dis(_d) {}
     
    inline bool operator < (const data &x) const {
        return dis < x.dis;
    }
};
 
priority_queue <data> h;
int n, k, K;
int now[Di];
 
inline int read() {
    int x = 0, sgn = 1;
    char ch = getchar();
    while (ch < '0' || '9' < ch) {
        if (ch == '-') sgn = -1;
        ch = getchar();
    }
    while ('0' <= ch && ch <= '9') {
        x = x * 10 + ch - '0';
        ch = getchar();
    }
    return sgn * x;
}
 
inline ll sqr(ll x) {
    return x * x;
}
 
inline ll dis(int p[Di], int q[Di]) {
    int i;
    ll res = 0;
    for (i = 0; i < k; ++i)
        res += sqr(p[i] - q[i]);
    return res;
}
 
#define P p -> p
#define Ls p -> s[0]
#define Rs p -> s[1]
inline void kd_new(kd_node *&p) {
    p = ++cnt_kd;
    memcpy(P, now, sizeof(P));
    Ls = Rs = null;
}
 
void kd_insert(kd_node *&p, int dep) {
    if (p == null) {
        kd_new(p);
        return;
    }
    bool d = P[dep] < now[dep];
    kd_insert(p -> s[d], (dep + 1) % k);
}
 
void kd_query(kd_node *p, int dep) {
    if (p == null) return;
    h.push(data(p, dis(P, now)));
    while ((int) h.size() > K) h.pop();
    bool d = P[dep] < now[dep];
    kd_query(p -> s[d], (dep + 1) % k);
    if ((int) h.size() < K || h.top().dis > sqr(now[dep] - P[dep]))
        kd_query(p -> s[!d], (dep + 1) % k);
}
 
void answer() {
    int i, j;
    while (!h.empty()) h.pop();
    kd_query(kd_root, 0);
    printf("the closest %d points are:\n", K);
    while ((int) h.size() > K) h.pop();
    i = K;
    while (i)
        ans[i--] = h.top().from, h.pop();
    for (i = 1; i <= K; ++i) {
        for (printf("%d", ans[i] -> p[0]), j = 1; j < k; ++j)
            printf(" %d", ans[i] -> p[j]);
        puts("");
    }
}
#undef P
#undef Ls
#undef Rs
 
int main() {
    int i, j, Q;
    null = mempool;
    null -> s[0] = null -> s[1] = null;
    while (scanf("%d%d", &n, &k) != EOF) {
        cnt_kd = mempool;
        kd_root = null;
        for (i = 1;  i <= n; ++i) {
            for (j = 0; j < k; ++j) now[j] = read();
            kd_insert(kd_root, 0);
        }
        Q = read();
        while (Q--) {
            for (j = 0; j < k; ++j) now[j] = read();
            K = read();
            answer();
        }
    }
    return 0;
}

(p.s. 感觉kd树写的不优美！！！要改！！！)