hdu 4347 The Closest M Points

Problem

acm.split.hdu.edu.cn/showproblem.php?pid=4347

vjudge.net/contest/187908#problem/C

Reference

详解KDTree

K-D树

Meaning

给出 k 维空间的 n 个点，q 个询问，每次给出一个点，问 n 个点中离它最近的 m 个点是哪些

Notes

我的理解，K-D树就是用来分割这个 k 维空间的，树上的每个结点都是一个分割点，以它为界，以某个维度为考量依据，把空间劈成两半。


（图片来自详解KDTree）

以二维平面为例，可以先按 x 轴方向来，在 n 个点中选个中点，把平
4000
面分成两半，这个中点就是树根（第1层），位于它两侧的点就分别位于它的两棵子树里。递归考虑它的子树（第2层），这次轮到考虑 y 轴方向，又是在这一侧的点里，按 y 来选个中点，以它为界又分两半…依此类推地递归处理。

上图就是先用黄点分左右、再用红点分上下、再用绿点分左右、再用蓝点分上下。

更高维的空间也类似，选取分割点的方法也有很多，这种简单版的就是各个维度顺序轮着用。

nth_element(first, nth, last, comp)

是 STL 里的模板函数，把处于区间 [ first，last）里的元素的第 n 个元素，放在第 n 个位置，程序中用它来找中点，并把它放到中间位置。像快排里的 partition。

Analysis

这道题要找离 p 点最近的 m 个点，大思路是首先把 K-D树建好后，从根到叶子一层层地搜 p 这个点被这棵树分到哪个区间，讲道理在这个区间附近找到的点应该离 p 点最近。搜完叶子再回溯，再看看隔壁子树有没有可能是更优的点。

Code

#include <cstdio>
#include <algorithm>
#include <queue>
#include <utility>
using namespace std;
const int N = 50000, K = 5;

int k, idx;

struct node
{
int v[K];
// 按照第 idx 个维度来排序
// 主要用于 nth_element()
bool operator < (const node &rhs) const
{
return v[idx] < rhs.v[idx];
}
} kdt[N<<2], p
, aim;

bool ex[N<<2]; // 标记树上的点是否存在

void build(int l, int r, int rt, int d)
{
// 当前结点存在
ex[rt] = true;
// 两个子结点标记为不存在先
ex[rt<<1] = ex[rt<<1|1] = false;
// 深度对维数求余
// 求出当前用哪个维度来分割空间
idx = d % k;
int m = l + r >> 1;
// 找区间的中点
// 并以它为分割点分割空间
// 比它小的放左子树，其它放右子树
nth_element(p + l, p + m, p + r + 1);
// 中点做分割点
// 也即当前的树上结点
kdt[rt] = p[m];
// [l , m - 1] 还有点 -> 左子树存在
if(l < m)
build(l, m - 1, rt<<1, d + 1);
// [m + 1 , r] 还有点 -> 右子树存在
if(r > m)
build(m + 1, r, rt<<1|1, d + 1);
}

inline int sqr(int x)
{
return x * x;
}

typedef pair<int,node> P;
// 用来存档前找到的候选点
// 把离目标点最远的候选点放在堆顶
// 当找到更近的点，就替换它
priority_queue<P> que;

void query(int rt, int m, int d)
{
if(!ex[rt])
return;
// 当前树结点（分割点）
P now(0, kdt[rt]);
// 它与与目标点的距离（的平方）
for(int i = 0; i < k; ++i)
now.first += sqr(aim.v[i] - kdt[rt].v[i]);

// 当前轮到考虑 dim 这一维
int dim = d % k;
// （如果有）优先搜索 one 子树，其次是 two 子树
int one = rt<<1, two = rt<<1|1;
// 看目标点被分在哪一侧
// 优先考虑那一侧
if(aim.v[dim] >= kdt[rt].v[dim])
swap(one, two);

// one 子树存在 -> 搜它
if(ex[one])
query(one, m, d + 1);
// 标记是否要搜 two 子树
bool fnd = false;
// 还不够 m 个点
if(que.size() < m)
{
que.push(now); // 直接添加
fnd = true; // 还要到 two 子树继续找更多的点
}
else // 已经够 m 个点 -> 考虑找更优的点替换
{
// 当前点更优 -> 替换
if(now.first < que.top().first)
{
que.pop();
que.push(now);
}
// two 子树有希望找到更优的点
// 搜隔壁子树
if(sqr(aim.v[dim] - kdt[rt].v[dim]) < que.top().first)
fnd = true;
}
// two 子树存在且决定要搜它
if(ex[two] && fnd)
query(two, m, d + 1);
}

int main()
{
int n;
while(scanf("%d%d", &n, &k) != EOF)
{
for(int i = 0; i < n; ++i)
for(int j = 0; j < k; ++j)
scanf("%d", &p[i].v[j]);

build(0, n - 1, 1, 0);

int t;
scanf("%d", &t);
for(int m; t--; )
{
for(int i = 0; i < k; ++i)
scanf("%d", &aim.v[i]);
scanf("%d", &m);
query(1, m, 0);
printf("the closest %d points are:\n", m);
for(int i = 0; !que.empty(); que.pop(), ++i)
p[i] = que.top().second;
for(int i = m - 1; ~i; --i)
for(int j = 0; j < k; ++j)
printf("%d%c", p[i].v[j], j + 1 == k ? '\n' : ' ');
}
}
return 0;
}

Post Script

讲真，这种写法的空间开得很谜，感觉有点奇怪，也有点浪费。

在写法上个人推荐看这一篇，感觉更更自然（像用指针动态开辟的写法，只是换成数组模拟指针），包含了最远点和最近点的估价。

而这一篇包含动态加入新点。