kd 树实现查找最近邻

在文章https://blog.csdn.net/App_12062011/article/details/51986805
的基础上修改的，原文章中的代码在回溯部分有不正确，回溯时如果要进入父节点的另一子树应该递归查找子树中的最近邻

#include <iostream>
#include <vector>
#include <cmath>
#include <algorithm>
#include <cstdlib>
using namespace std;
struct kdTree{
 vector<double> data;
 kdTree* parent;
 kdTree* leftChild;
 kdTree* rightChild;
 //默认构造函数 
 kdTree(){parent = leftChild = rightChild = NULL;}
 //判断kd树是否为空
 bool isEmpty()
 {
  return data.empty();
 }
 //判断kd树是否只是一个叶子结点
 bool isLeaf()
 {
  return (!data.empty())&&rightChild == NULL && leftChild == NULL;
 } 
 //判断kd树是否是根结点
 bool isRoot()
 {
  return (!data.empty())&&parent == NULL;
 }
 //判断kd树是否为左子树
 bool isLeftChild()
 {
  return parent->leftChild->data == data;
 }
 //判断kd树是否为右子树
 bool isRightChild()
 {
  return parent->rightChild->data == data;
 }
   
}; 
template<typename T>
vector<vector<T> > Transpose(vector<vector<T> > Matrix)
{
 unsigned row = Matrix.size();
 unsigned col = Matrix[0].size();
 vector<vector<T> > Trans(col,vector<double>(row,0));
 for(unsigned i=0;i<col;i++){
  for(unsigned j=0;j<row;j++)
   Trans[i][j] = Matrix[j][i];
 }
 return Trans;
}

template<typename T>
T findMiddleValue(vector<T> vec)
{
 sort(vec.begin(),vec.end());
 int pos = vec.size()/2;
 return vec[pos];
}

//构建kd树
void buildkdTree(kdTree *tree,vector<vector<double> >data,unsigned depth)
{
 //样本的数量
 unsigned sampleNum = data.size();
 //终止条件
 if(sampleNum == 0)
 {
  return;
 }
 if(sampleNum == 1)
 {
  tree->data = data[0];
  return;
 }
 
 //样本的维度
 unsigned k = data[0].size();
 vector<vector<double> > transData = Transpose(data);
 //选择切分属性
 unsigned splitAttribute = depth % k;
 vector<double> splitAttributeValues = transData[splitAttribute];
 //选择切分值
 double splitValue = findMiddleValue(splitAttributeValues);
 //根据选定的切分属性和切分值将数据分为两个集合
 vector<vector<double> > leftPart;
 vector<vector<double> > rightPart;
 for(unsigned i=0;i<sampleNum;i++)
 {
  if(splitAttributeValues[i]==splitValue && tree->data.empty())
   tree->data = data[i];
  else{
   if(splitAttributeValues[i] < splitValue)
    leftPart.push_back(data[i]);
   else
    rightPart.push_back(data[i]);
  }
 }
 //子集递归调用buildkdTree函数
 if(!leftPart.empty())
 {
  tree->leftChild = new kdTree;
  tree->leftChild->parent = tree;
  buildkdTree(tree->leftChild,leftPart,depth + 1);
 }
 if(!rightPart.empty())
 {
  tree->rightChild = new kdTree;
  tree->rightChild->parent = tree;
  buildkdTree(tree->rightChild,rightPart,depth + 1);
 }
 
 } 
 //逐层打印kd树
void printkdTree(kdTree *tree,unsigned depth)
 {
  for(unsigned i=0;i<depth;i++)
   cout<<"\t";
  for(unsigned j=0;j<tree->data.size();j++)
   cout<<tree->data[j]<<",";
  cout<<endl;
  if(tree->leftChild == NULL && tree->rightChild == NULL)
   return;
  else//非叶子节点
  {
   if(tree->leftChild != NULL)
   {
    for(unsigned i = 0;i<depth + 1;i++)
     cout << "\t";
    cout << "left:";
    printkdTree(tree->leftChild,depth + 1);
    } 
    
    cout << endl;
    if(tree->rightChild != NULL)
    {
     for(unsigned i = 0;i < depth + 1;i++)
      cout << "\t";
     cout << "right:";
     printkdTree(tree->rightChild,depth + 1);
    }
    cout << endl;
 }   
  
 } 
//计算空间中两个点的距离
 double measureDistance(vector<double> point1,vector<double> point2, unsigned method)
 {
  if(point1.size() != point2.size())
  {
   cerr << "Dimensions don't match !!" ;
   exit(1);
  }
  double res = 0;
  switch(method)
  {
   case 1://曼哈顿距离 
      {
       for(unsigned i=0;i<point1.size();i++)
     res += abs(point1[i]-point2[i]);
     return res;
      }  
   case 2://欧式距离 
    {
     for(unsigned i=0;i < point1.size();i++)
     {
      res += pow((point1[i] - point2[i]),2);
     }
     return sqrt(res);
    }
     default:
      {
       cerr << "Invalid method!!" << endl;
       return -1;
   }
  }
  }
//在kd树tree中搜索目标点goal的最近邻
  //输入：目标点；已构造 的kd树
  //输出：目标点的最近邻 
  kdTree* searchNearestNeighbor(vector<double> goal,kdTree *tree)
  {
   /*step1：在kd树中找出包含目标点的叶子节点：从根节点出发，
   递归的向下访问kd树 ，若目标点的当前维的坐标小于切分点的坐标，
   则移动到左子树中查找，若否则移动到右子树中查找，
   直到子结点为叶子结点为止，以此叶子结点为 “当前最近点”
 */
  unsigned k = tree->data.size();//计算出数据的维度 
  unsigned d = 0; //维度 初始化为0，即从第1维开始
  kdTree *currentTree = tree;
  while(!currentTree->isLeaf())
  {
   unsigned index = d%k;//计算当前维
  if(currentTree->rightChild == NULL|| goal[index] < currentTree->data[index])  
  {
   currentTree = currentTree->leftChild;
   
  }
  else
  {
   currentTree = currentTree->rightChild;
  }
  d++;
  }
  kdTree *nearest_point = currentTree;
  
  
  /*step2: 递归地向上回退，在每个结点进行如下操作：
  (a)如果该结点保存的实例比当前最近点距离目标点更近，则以该例点为“当前最近点”
  (b)当前最近点一定存在于某结点一个子结点对应的区域，检测该子节点的父节点的另一子节点对应
  区域是否有更近的点（即检查另一子节点对应的区域是否与以目标点为球心，以目标点与“当前最近点”
  间的 距离为半径的球体相交）；如果相交，可能在另一个子节点对应的区域内存在距目标点更近的点，
  移动到另 一个子节点，接着递归进行最近邻搜索，如果不想交，向上回退*/
  
  double currentDistance = measureDistance(goal,nearest_point->data,2);
  while(currentTree != tree)
  {
   double districtDistance = abs(goal[(--d)%k] - currentTree->parent->data[(--d)%k]); 
   if(districtDistance < currentDistance) //是否进入父节点的另一区域搜索 
   {//如果当前子kd树的根结点是其父结点的左孩子，则搜索其父节点的右孩子节点所代表的区域，反之亦然
     kdTree * possible_point = NULL;
    if(currentTree->isLeftChild()){
       if(currentTree->parent->rightChild != NULL)
      possible_point = searchNearestNeighbor(goal,currentTree->parent->rightChild); 
   
    }
   else if(currentTree->parent->leftChild != NULL)
    possible_point = searchNearestNeighbor(goal,currentTree->parent->leftChild); 
    
   if(possible_point != NULL)
   {
    double distance = measureDistance(goal,possible_point->data,2);
    if(distance < currentDistance)
    {
     currentDistance = distance;
     nearest_point = possible_point;
    }
   }
   
   }
   double parentDistance = measureDistance(goal,currentTree->parent->data,2);
   if(parentDistance < currentDistance)
   {
    currentDistance = parentDistance;
    nearest_point = currentTree->parent;
   }
   currentTree = currentTree->parent;
  }
  
  return nearest_point;
  }

int main()
{
int sampleNum = 10;
   int dimension = 5;
   vector<double>goal(dimension,0);
   for(unsigned i=0;i<dimension;i++)
    goal[i] = rand()%10;
   vector<vector<double> > train(sampleNum,vector<double>(dimension,0));
   double minDistance = 10;
   vector<double> closest;
   for(unsigned i = 0;i < sampleNum; i++){
    for(unsigned j = 0;j < dimension; j++)
     train[i][j] = rand()%10;
    double distance = measureDistance(goal,train[i],2);
    if(distance < minDistance)
    {
     closest = train[i];
     minDistance = distance;
    }
   }
kdTree* tree = new kdTree;
   buildkdTree(tree,train,0);
   printkdTree(tree,0);
   
   kdTree* nearestNeighbor = searchNearestNeighbor(goal,tree);
   vector<double>::iterator beg = nearestNeighbor->data.begin() ;
   cout << "The goal is:";
   for(unsigned i=0;i<dimension;i++)
  cout<<goal[i]<<",";
 cout<<endl;
   cout << "The nearest neighbor is:";
   while(beg != nearestNeighbor->data.end()) cout << *beg++ << ",";
 cout<<endl;
 for(unsigned i=0;i<dimension;i++)
  cout<<closest[i]<<",";
 cout<<endl;
 return 0;

}