之前写过一篇使用matlab建立Kd-tree并进行k-NN查询。之前写完之后没怎么测试就直接丢上来了,这些天有人后台私信问相关的问题后才发现搜索算法的实现好像有一些问题,而且整体代码比较杂乱,想着重新实现一遍,因此有了这篇。
在网上看了很多人的讲解,具体实现的方法千差万别,有的只在叶节点存放真实数据,在根节点中存放分割信息。有的所有节点都能存放数据信息。有的为了回溯时减少递归次数在节点中存放自身所代表的矩形框大小。在这里我也只实现一种我觉得方便的实现方法,就是每一个节点都存放一个数据点以及基于该数据点的分割信息,每次分割选择方差最大的维度进行分割,另外用绘制函数把2维下的kd-tree绘制出来,这样也更方便可视化了解算法思想。
用matlab来实现这种依赖于指针的树结构老实说并不方便,迫不得已使用数组的索引来代替指针来实现子节点到父节点回溯。后续将自己实现的kd-tree搜索和暴力搜索进行比较,在数据量较小或者数据维度较多时,kd-tree的搜索时间比暴力搜索还要长好多。而且严重依赖于构造出的kd-tree的质量,同样的数据量和数据维度,在几次不同的测试下kd-tree的最近邻搜索耗时可能相差数倍。只是在数据维度小于6维,并且数据量大小大于 2 14 2^{14} 214时,kd-tree的最近邻搜索才开始表现出优势。这很可能是与我的实现方法有关,在matlab上实现主要是了解算法的思想,以及可视化绘图方便,可能对于真实算法性能比较方面只能定性分析,没法准确定量。
kd-tree的最近邻搜索的实现主要参考了李小文的知乎专栏文章:KD Tree的原理及Python实现
运行结果如图:
具体代码如下:
clc
clear
close all
% 数据维度
% 建立kd-tree和最近邻搜索对于维度没限制,只是绘制kd-tree只有数据为2维才有效
DATA_DIM = 2;
% 数据量大小
DATA_SIZE = 20;
data = rand(DATA_SIZE, DATA_DIM);
data = unique(data, 'rows');
% 建立kd-tree前应当归一化数据,保证以方差作为选取划分维度方式的合理性
% 这里每个维度都是0到1的随机数因此不用再归一化
kd_tree = kd_tree_build(data);
if DATA_DIM == 2
figure;
hold on
Limit_X = [0, 1];
Limit_Y = [0, 1];
plot_kd_tree(gca, kd_tree, Limit_X, Limit_Y);
end
% 目标点
target = rand(1, DATA_DIM);
if DATA_DIM == 2
plot(target(1), target(2), '*');
end
% 比较kd-tree搜索和暴力搜索的用时
time_kd = 0;
time_simple = 0;
for i = 1: 10000
tic
res_kd = kd_tree_search(kd_tree, target);
elapsedTime = toc;
time_kd = time_kd + elapsedTime;
tic
res_simple = simple_search(data, target);
elapsedTime = toc;
time_simple = time_simple + elapsedTime;
assert(sqrt(sum((res_kd - res_simple).^2)) < 0.0000001, 'error:kd-tree搜索和暴力搜索结果不一致');
end
disp(['kd-tree搜索用时' num2str(time_kd) '秒']