题意:给你n个k维的点,再给你一个目标点x,让你查询离x最近的M个点?
分析:【KD树模板】
#include<cstdio>
#include<bits/stdc++.h>
#include<algorithm>
#include<string.h>
using namespace std;
const int N = 50000; //最大点数
short currentDim;
#define square(x) (x)*(x)
int n; //n个点
short K, t; //K维,t次询问
struct point {
short pos[5];
bool flag;
point() {
flag = false;
}
bool operator < (const point &a) const {
return pos[currentDim] < a.pos[currentDim];
}
};
point input_v[N];
priority_queue<pair<double, point>> pq;
point kd_tree[4*N];
inline void build(int n, short K, int p, int r, int index, short depth) {
if(p <= r) {
currentDim = depth % K;
int mid = (p + r) / 2;
nth_element(input_v+p, input_v+mid, input_v+r+1);
kd_tree[index] = input_v[mid];
kd_tree[index].flag = true;
build(n, K, p, mid-1, 2*index, depth+1);
build(n, K, mid+1, r, 2*index+1, depth+1);
}
}
inline void query(point p, short closestM, short K, int index, short depth) {
if(!kd_tree[index].flag) return;
pair<double, point> current_node(0, kd_tree[index]);
for(int i = 0; i < K; i++) {
current_node.first += square(p.pos[i]-current_node.second.pos[i]);
}
short idm = depth % K;
short flag = 0;
int lchild = 2*index;
int rchild = 2*index+1;
if(kd_tree[index].pos[idm] <= p.pos[idm]) {
query(p, closestM, K, rchild, depth+1);
} else {
query(p, closestM, K, lchild, depth+1);
}
if(pq.size() < closestM) {
pq.push(current_node);
flag = 1;
} else {
if(current_node.first < pq.top().first) {
pq.pop();
pq.push(current_node);
}
if(square(p.pos[idm] - kd_tree[index].pos[idm]) < pq.top().first) {
flag = 1;
}
}
if(flag) {
if(kd_tree[index].pos[idm] <= p.pos[idm])
query(p, closestM, K, lchild, depth+1);
else
query(p, closestM, K, rchild, depth+1);
}
}
int main() {
while(scanf("%d%hd",&n, &K) != EOF) {
for(int i = 0; i < n; i++) { //n个点
for(short j = 0; j < K; j++) {
scanf("%hd", &input_v[i].pos[j]);
}
}
memset(kd_tree, 0, sizeof(point)*4*N);
build(n, K, 0, n-1, 1, 0); //建树
scanf("%hd", &t);
for(short c = 0; c < t; c++) { //t次询问
short closestM;
point tmpv;
for(short j = 0; j < K; j++) { //K维
scanf("%hd", &tmpv.pos[j]);
}
scanf("%hd", &closestM); //最近M个点
query(tmpv, closestM, K, 1, 0); //查询
printf("the closest %hd points are:\n", closestM);
point pt[10];
for(short i = 0; !pq.empty(); i++) {
pt[i] = pq.top().second; //最近M个点保存到pq队列
pq.pop();
}
for(short i = closestM-1; i >= 0; i--) {
for(short j = 0; j < K; j++) {
printf("%hd", pt[i].pos[j]);
if(j == K-1) printf("\n");
else printf(" ");
}
}
}
}
return 0;
}