KD-Tree构建及邻近点搜索 C++实现

KD-Tree构建及邻近点搜索 C++实现

1、KD-Tree构建原理

KD-Tree是一种树形数据结构，常用于点云处理。KD-Tree会在其非叶结点上对数据进行划分，在叶子结点上存储数据。

结点数据结构如下：

//KD树结点
struct TreeNode
{
	//数据维度，如三维点，二维点
	int dim = 1;

	//当前结点的分割维度，表示在那个维度进行分割
	int split_dim = -1;

	//当前结点所有叶子结点存储数据的大小
	int data_size = 0;

	//是否为左叶子结点，false表示当前结点是父结点的右结点
	bool is_left = true;

	//孩子结点
	TreeNode* chirld[2];

	//父结点
	TreeNode* father;

	//当前结点所有数据的平均值
	double* split_center;

	//当前结点所存储数据的索引，非叶子结点会在数据分割后清空内存
	vector<int> data;

	//有参构造，in_dim为输入的维度
	TreeNode(int in_dim)
	{
		dim = in_dim;
		chirld[0] = nullptr;
		chirld[1] = nullptr;
		father = nullptr;
		split_center = new double[dim];
		memset(split_center, 0, dim * sizeof(double));
	}
};

通常情况下，KD-Tree的构建过程有三步：

（1）计算最佳的划分维度以及当前待划分数据的平均值，通常以当前结点待划分数据的最大方差维度作为划分轴。

（2）根据划分轴及数据平均值对数据进行划分。

（3）递归左孩子，右孩子，直到所有数据划分完成。

当数据量特别大时，使用递归创建的方式将会使存储函数的栈空间都被使用完，后续在创建函数会弹出异常。

void KDTreeSearch::GetDataFeatrue(
	const vector<vector<double>>& data, 
	const vector<int>&            index, 
	int                           dim, 
	double*                       center, 
	int&                          split_dim)
{
	const int data_size = index.size();
	
	vector<double> average(dim, 0.);
	vector<double> standard(dim, 0.);

	for (int i = 0; i < data_size; i++)
	{
		int k = index[i];
		const auto& p = data[k];
		for (int j = 0; j < dim; j++)
		{
			average[j] += p[j];
		}
	}
	for (int j = 0; j < dim; j++)
	{
		average[j] /= data_size;
	}
	memcpy(center, average.data(), dim * sizeof(double));

	for (int i = 0; i < data_size; i++)
	{
		int k = index[i];
		const auto& p = data[k];
		for (int j = 0; j < dim; j++)
		{
			standard[j] += pow(p[j] - average[j], 2);
		}
	}
	for (int j = 0; j < dim; j++)
	{
		standard[j] = sqrt(standard[j] / data_size);
	}

	split_dim = 0;
	double max_val = standard[split_dim];
	for (int i = 1; i < dim; i++)
	{
		if (max_val < standard[i])
		{
			max_val = standard[i];
			split_dim = i;
		}
	}
}

void KDTreeSearch::CreateKDTree(const vector<vector<double>>& data, TreeNode* head)
{
	//当前结点只有一个数据时，退出
	if (head->data_size <= 1) return;

	//1.获取分割维度以及平均值
	GetDataFeatrue(data, head->data, head->dim, head->split_center, head->split_dim);

	//2.构造左右子树
	for (int i = 0; i < 2; i++)
	{
		head->chirld[i] = new TreeNode(head->dim);
		head->chirld[i]->is_left = 1 - i;            //表示当前结点是否为父结点的左孩子
		head->chirld[i]->father = head;              //将孩子结点指向父结点
	}

	//3.根据分割维度对数据进行分割
	for (int i = 0; i < head->data_size; i++)
	{
		int k = head->data[i];
		const auto& p = data[k];
		if (p[head->split_dim] < head->split_center[head->split_dim])
		{
			//小于分割值，加入左孩子
			head->chirld[0]->data.push_back(k);
			head->chirld[0]->data_size++;
		}
		else
		{
			//大于分割值，加入右孩子
			head->chirld[1]->data.push_back(k);
			head->chirld[1]->data_size++;
		}
	}

	//4.清空当前结点存储的数据
	vector<int>().swap(head->data);

	//5.递归
	for (int i = 0; i < 2; i++)
	{
		CreateKDTree(data, head->chirld[i]);
	}
}

2、KD-Tree的遍历与删除

KD-Tree的遍历与二叉树的遍历相同，可以使用先序、中序、后序等遍历方式，但KD-Tree只有叶子结点存储真正的数据，所以非叶结点直接跳过即可。下面为KD-Tree的后续遍历方式。

void KDTreeSearch::ErgodicKDTree(TreeNode* head, vector<int>& index)
{
	if (head == nullptr) return;
	//后续遍历叶子结点所有元素
	ErgodicKDTree(head->chirld[0], index);
	ErgodicKDTree(head->chirld[1], index);
	if (head->chirld[0] == nullptr)
		index.insert(index.end(), head->data.begin(), head->data.end());
}
<