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;
}

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