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

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