KD树(K-dimension tree),个人理解就是一棵扩展到k维空间上的二叉搜索树,它在每一层用不同的维度进行二分,将空间中的点按一定的顺序组织成二叉树的形式,以方便搜索,具体的内容就不在这里细述了。
来看看这道题吧,题目大意:给定K(<=5)维空间中的n(<=5000)个点,每个坐标的绝对值不超过10000,做t次(<=10000)查询,询问距离指定坐标最近的m个点(m<=10)
直接暴搜的话肯定会超时,这里用KD Tree存储这些点,然后剪枝搜索
这道题有个大坑,数据是多组的,“There are multiple test cases. Process to end of file.”,样例又没有体现,对我这种不认真读题的真是噩梦一般.......
先来思考这么一个问题,如何查找距离某坐标最近的节点?可以根据每一层的划分依据进入左子树或右子树查找,维护一个最小距离,这样当遍历到叶节点的时候我们就能得到当前搜索范围中的最近点,当然这不一定是最近的节点,之所以优先遍历这一条路径是为了引发更多的剪枝。在一层层的向上返回时再去进入另一颗子树查找
那么如何剪枝呢?对于某节点所处的一层来说,这个节点在这个维度上将空间二分了,可以看成一条数轴,如果以目标点为中点、当前最小距离为半径的区间包不包含含该节点,那么显然最短距离只可能落在目标节点一层,而这一侧一定是递归搜索时先进入的子树,再返回到该节点时,发现条件不满足,那么就没有必要进行另一棵子树的搜索了。
这样就完成了找到最近的一个点的任务,那么如何找到最近的M个点?这个就简单了,上面的问题中只要保存找到的最近点就行了,现在保存当前找到的最近M个点的集合,维护最小距离时维护的是当前搜索的节点中第M短的长度,如果集合中元素不足M个,就直接将点加入集合中。具体实现上可以用最大堆来维护,不过因为这道题M<=10,普通的插入排序效率也够了
思想就是这些,下面是代码(还是老样子,写的比较难看....不知为啥一时脑子进水用了索引排序==)
#include <iostream>
#include <cstdio>
#include <cmath>
#define MAXN 50010
#define SQ(X) (X)*(X)
using namespace std;
int data[MAXN][12]; // 点数据
int index[MAXN]; // 点的索引,用于分割
int ans[MAXN]; // 存放查找到的点
struct treeNode
{
int value; // 代表点的索引
int splitIndex; // 分割的维度
treeNode * left;
treeNode * right;
// treeNode(int _val, int _sIndex): value(_val), splitIndex(_sIndex), left(NULL), right(NULL) {}
}trNode[MAXN];
struct KDtree
{
int dimension;
int cnt; // 树中的节点
treeNode * root; // 根节点
treeNode * treeArray[15]; // 保存当前维护的最近m个节点
int dist[15]; // 当前最近m个节点的距离
int top; // 当前维护的集合中的节点个数
KDtree(int k): dimension(k)
{
cnt = 0;
top = 0;
}
int partition(int * dataIndex, int s, int e, int splitIndex) // 将data中的数据根据splitIndex,做一次分割,返回分割归位的位置
{
int tmp = dataIndex[e - 1];
int i = s, j = e - 1;
while (i < j)
{
//printf("%d %d\n", i, j);
while (i < j && data[dataIndex[i]][splitIndex] <= data[tmp][splitIndex])
i++;
if (i == j)
break;
dataIndex[j] = dataIndex[i];
j--;
while (i < j && data[dataIndex[j]][splitIndex] >= data[tmp][splitIndex])
j--;
if (i == j)
break;
dataIndex[i] = dataIndex[j];
i++;
}
dataIndex[i] = tmp;
return i;
}
treeNode * build(int * dataIndex, int s, int e, int d) // 递归建树
{
//printf("%d %d\n", s, e);
if (s >= e)
return NULL;
int splitIndex = d % dimension; // 据说选择划分的维度根据方差效果是最优的,偷懒了> <
int mid = (s + e) >> 1, tmp = e, ptCnt = 0; // 建树的过程中实际上要找到中位数,可以用STL中的nth_element,我这个写的貌似是卡特殊数据了,贡献了无数次TLE....不过对于这道题来说任意的划分一次貌似效果也不错
/* while (tmp != mid && ptCnt < 5)
{
//printf("%d\n", tmp);
if (tmp > mid)
tmp = partition(dataIndex, s, tmp, splitIndex);
else
tmp = partition(dataIndex, tmp + 1, e, splitIndex);
ptCnt++;
}*/
tmp = partition(dataIndex, s, e, splitIndex); // 任意分割一次
trNode[cnt].value = dataIndex[tmp];
trNode[cnt].splitIndex = splitIndex;
trNode[cnt].left = NULL;
trNode[cnt].right = NULL;
treeNode * tmpTreeNode = &trNode[cnt++];
//treeNode * tmpTreeNode = new treeNode(dataIndex[tmp], splitIndex);
tmpTreeNode->left = build(dataIndex, s, tmp, d + 1);
tmpTreeNode->right = build(dataIndex, tmp + 1, e, d + 1);
return tmpTreeNode;
}
void buildKD(int n, int * dataIndex) // 根据索引数组构建KDtree
{
root = build(dataIndex, 0, n, 0);
}
int distance(int * a, int * b) // 返回两个点之间的距离
{
int ans = 0;
for (int i = 0; i < dimension; ++i)
{
ans += (a[i] - b[i]) * (a[i] - b[i]);
}
return ans;
}
void find(int * pt, treeNode * root, int m) // 递归查找节点
{
if (root == NULL)
return;
int splitIndex = root->splitIndex;
treeNode * another; // 未选择的分支
int tmpDist = distance(pt, data[root->value]);
int p = top; // 维护当前距离最近的m个点的数组
while (p > 0 && tmpDist < dist[p - 1]) // 插入当前节点
{
treeArray[p] = treeArray[p - 1];
dist[p] = dist[p - 1];
p--;
}
treeArray[p] = root;
dist[p] = tmpDist;
if (top < m)
++top;
if (pt[splitIndex] < data[root->value][splitIndex]) // 递归搜索
{
another = root->right;
find(pt, root->left, m);
}
else
{
another = root->left;
find(pt, root->right, m);
}
if (top < m || SQ(pt[splitIndex] - data[root->value][splitIndex]) < dist[top - 1]) // 如果另一个分支可能存在最优解继续搜索
find(pt, another, m);
}
void findML(int * pt, int m, int * ansArray)
{
top = 0;
find(pt, root, m);
for (int i = 0; i < m; ++i)
{
ansArray[i] = treeArray[i]->value;
}
return;
}
};
int main()
{
int n, m, k, t;
int pt[16];
while (scanf("%d%d", &n, &k) == 2)
{
KDtree kd(k);
for (int i = 0; i < n; ++i)
{
for (int j = 0; j < k; ++j)
{
scanf("%d", data[i] + j);
}
index[i] = i;
}
kd.buildKD(n, index);
scanf("%d", &t);
while (t--)
{
for (int i = 0; i < k; ++i)
{
scanf("%d", pt + i);
}
scanf("%d", &m);
kd.findML(pt, m, ans);
printf("the closest %d points are:\n", m);
for (int i = 0; i < m; ++i)
{
int tmpIndex = ans[i];
printf("%d", data[tmpIndex][0]);
for (int j = 1; j < k; ++j)
{
printf(" %d", data[tmpIndex][j]);
}
printf("\n");
}
}
}
return 0;
}