裸的KD-tree，还是在估计要不要进入子树的时候判断一下就好了，剩下都一样
判断的方法就是看现在答案个数是否小于k，和答案是否会经过分割线。
1 /************************************************************** 2 Problem: 3053 3 User: rausen 4 Language: C++ 5 Result: Accepted 6 Time:2164 ms 7 Memory:3572 kb 8 ****************************************************************/ 9 10 #include <cstdio> 11 #include <cstring> 12 #include <algorithm> 13 #include <queue> 14 15 using namespace std; 16 typedef long long ll; 17 18 const int Cnt_kd = 1e5 + 5; 19 const int Di = 5; 20 21 struct kd_node { 22 kd_node *s[2]; 23 int p[Di]; 24 } *kd_root, mempool[Cnt_kd], *cnt_kd, *null, *ans[25]; 25 26 struct data { 27 kd_node *from; 28 ll dis; 29 data() {} 30 data(kd_node *_f, ll _d) : from(_f), dis(_d) {} 31 32 inline bool operator < (const data &x) const { 33 return dis < x.dis; 34 } 35 }; 36 37 priority_queue <data> h; 38 int n, k, K; 39 int now[Di]; 40 41 inline int read() { 42 int x = 0, sgn = 1; 43 char ch = getchar(); 44 while (ch < '0' || '9' < ch) { 45 if (ch == '-') sgn = -1; 46 ch = getchar(); 47 } 48 while ('0' <= ch && ch <= '9') { 49 x = x * 10 + ch - '0'; 50 ch = getchar(); 51 } 52 return sgn * x; 53 } 54 55 inline ll sqr(ll x) { 56 return x * x; 57 } 58 59 inline ll dis(int p[Di], int q[Di]) { 60 int i; 61 ll res = 0; 62 for (i = 0; i < k; ++i) 63 res += sqr(p[i] - q[i]); 64 return res; 65 } 66 67 #define P p -> p 68 #define Ls p -> s[0] 69 #define Rs p -> s[1] 70 inline void kd_new(kd_node *&p) { 71 p = ++cnt_kd; 72 memcpy(P, now, sizeof(P)); 73 Ls = Rs = null; 74 } 75 76 void kd_insert(kd_node *&p, int dep) { 77 if (p == null) { 78 kd_new(p); 79 return; 80 } 81 bool d = P[dep] < now[dep]; 82 kd_insert(p -> s[d], (dep + 1) % k); 83 } 84 85 void kd_query(kd_node *p, int dep) { 86 if (p == null) return; 87 h.push(data(p, dis(P, now))); 88 while ((int) h.size() > K) h.pop(); 89 bool d = P[dep] < now[dep]; 90 kd_query(p -> s[d], (dep + 1) % k); 91 if ((int) h.size() < K || h.top().dis > sqr(now[dep] - P[dep])) 92 kd_query(p -> s[!d], (dep + 1) % k); 93 } 94 95 void answer() { 96 int i, j; 97 while (!h.empty()) h.pop(); 98 kd_query(kd_root, 0); 99 printf("the closest %d points are: ", K); 100 while ((int) h.size() > K) h.pop(); 101 i = K; 102 while (i) 103 ans[i--] = h.top().from, h.pop(); 104 for (i = 1; i <= K; ++i) { 105 for (printf("%d", ans[i] -> p[0]), j = 1; j < k; ++j) 106 printf(" %d", ans[i] -> p[j]); 107 puts(""); 108 } 109 } 110 #undef P 111 #undef Ls 112 #undef Rs 113 114 int main() { 115 int i, j, Q; 116 null = mempool; 117 null -> s[0] = null -> s[1] = null; 118 while (scanf("%d%d", &n, &k) != EOF) { 119 cnt_kd = mempool; 120 kd_root = null; 121 for (i = 1; i <= n; ++i) { 122 for (j = 0; j < k; ++j) now[j] = read(); 123 kd_insert(kd_root, 0); 124 } 125 Q = read(); 126 while (Q--) { 127 for (j = 0; j < k; ++j) now[j] = read(); 128 K = read(); 129 answer(); 130 } 131 } 132 return 0; 133 }
(p.s. 感觉kd树写的不优美！！！要改！！！)