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 个点中选个中点，把平面分成两半，这个中点就是树根（第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[N], 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

讲真，这种写法的空间开得很谜，感觉有点奇怪，也有点浪费。
在写法上个人推荐看这一篇，感觉更更自然（像用指针动态开辟的写法，只是换成数组模拟指针），包含了最远点和最近点的估价。
而这一篇包含动态加入新点。