KDTree

学习链接：http://www.cnblogs.com/eyeszjwang/articles/2429382.html

下面实现的kdtree支持以下操作：
(1) 插入一个节点
(2) 插入n个节点
(3) 查找距离某个给定点距离最小的K个节点。

插入的时候可能会导致树严重不平衡，这个时候会重建某个子树。

 #include <algorithm>
 #include <queue>

 /***
   _NodeType: 节点类型
   _CompareFuncType: 比较函数类型,它应该接受两个_NodeType并返回0或者1
   _UpdateFuncType:  更新函数类型,它有三个参数，分别表示Father,leftSon,rightSon,参数的类型的 _NodeType*
   _DIMENSION:  维数
   _DISTANCE_TYPE: 两个_NodeType之间距离的类型
   _REBUILD_ALPHA: 这个参数用来维持树大致平衡，它应该大于50小于等于100
 ***/
 template<class _NodeType,
          class _CompareFuncType,
          class _UpdateFuncType,
          int _DIMENSION,
          class _DISTANCE_TYPE,
          int _REBUILD_ALPHA=>
 class KDTree
 {
 private:
     struct TreeNode
     {
         TreeNode(const _NodeType& _Node):m_Left(nullptr),m_Right(nullptr),m_Size(),
             m_Node(_Node) {}
         TreeNode() {}

         _NodeType m_Node;
         TreeNode* m_Left;
         TreeNode* m_Right;
         unsigned int m_Size;
     };

     TreeNode* NewNode(const _NodeType& _Node)
     {
         return new TreeNode(_Node);
     }

 public:
     KDTree(_CompareFuncType** _CompareFuncs,_UpdateFuncType* _UpdateFunc):m_Root(nullptr)
     {
         for(unsigned int Idx=;Idx<_DIMENSION;++Idx)
         {
             m_CompareFuncs[Idx]=_CompareFuncs[Idx];
         }
         m_UpdateFunc=_UpdateFunc;
     }

     void insert(const _NodeType& _Value)
     {
         TreeNode* BadTreeNode=nullptr;
         TreeNode* BadTreeNodeParent=nullptr;
         unsigned int BadTreeNodeDimension=;
         m_Root=Insert(m_Root,NewNode(_Value),,
             BadTreeNode,BadTreeNodeParent,BadTreeNodeDimension);

         if(BadTreeNode!=nullptr)
         {
             if(BadTreeNodeParent==nullptr)
             {
                 m_Root=RebuildTree(BadTreeNode,BadTreeNodeDimension);
             }
             else
             {
                 BadTreeNodeParent=RebuildTree(BadTreeNode,BadTreeNodeDimension);
             }
         }
     }

     template<class _NodeTypeBegin>
     void insert(_NodeTypeBegin _ValueArray,const unsigned int _ValueArraySize)
     {
         if(_ValueArraySize<=) return;
         if(m_Root==nullptr)
         {
             m_Root=BuildGroup(_ValueArray,,_ValueArraySize-,);
         }
         else
         {
             if(m_Root->m_Size<=(unsigned int)(_ValueArraySize*_REBUILD_ALPHA/))
             {
                 _NodeType* NewValueArray=new _NodeType[_ValueArraySize+m_Root->m_Size];
                 unsigned int NewValueArraySize=StoreSubtreeNodeIntoArray(m_Root,NewValueArray);
                 ClearSubTrees(m_Root);
                 for(unsigned int Idx=;Idx<_ValueArraySize;++Idx)
                 {
                     NewValueArray[NewValueArraySize++]=_ValueArray[Idx];
                 }
                 m_Root=BuildGroup(NewValueArray,,NewValueArraySize-,);
                 delete[] NewValueArray;
             }
             else
             {
                 for(unsigned int Idx=;Idx<_ValueArraySize;++Idx)
                 {
                     insert(_ValueArray[Idx]);
                 }
             }
         }
     }

     /***
       查找距离_Value "最近" 的_SearchNumber个元素 存储在_StoreAnswerArray

       _ComputeMinDistanceFunc接受两个_NodeType(first,second) 用来计算second范围内所有点到first的 "最近"距离
              _DISTANCE_TYPE是它的返回值类型
       _ComputeDistanceFunc接受两个_NodeType(first,second) 计算first和second的距离
              _DISTANCE_TYPE是它的返回值类型
       _CompareDistanceFunc它接受两个_DISTANCE_TYPE(first,second) 并返回0或者1
              1表示first小于second
       函数返回查找到的元素个数(有可能小于_SearchNumber)
     ***/
     template<class _ComputeDistanceFuncType,
              class _CompareDistanceFuncType>
     unsigned int searchKNear(const _NodeType& _Value,_NodeType* _StoreAnswerArray,
         const unsigned int _SearchNumber,_CompareDistanceFuncType* _CompareDistanceFunc,
         _ComputeDistanceFuncType* _ComputeMinDistanceFunc,
         _ComputeDistanceFuncType* _ComputeDistanceFunc)
     {
         if(_SearchNumber==) return ;
         unsigned int AnswerArrayElementNumber=;
         SearchKNear(m_Root,_Value,_StoreAnswerArray,AnswerArrayElementNumber,_SearchNumber,
             _CompareDistanceFunc,_ComputeMinDistanceFunc,_ComputeDistanceFunc);
         return AnswerArrayElementNumber;
     }

     unsigned int size() const
     {
         if(m_Root) return m_Root->m_Size;
         return ;
     }

 private:
     template<class _ComputeDistanceFuncType,
              class _CompareDistanceFuncType>
     void SearchKNear(
            TreeNode* _Root,const _NodeType& _Value,_NodeType* _StoreAnswerArray,
            unsigned int &_CurAnswerArrayElementNumber,
            const unsigned int _SearchNumber,_CompareDistanceFuncType* _CompareDistanceFunc,
            _ComputeDistanceFuncType* _ComputeMinDistanceFunc,
            _ComputeDistanceFuncType* _ComputeDistanceFunc)
     {
         if(_Root==nullptr) return;
         if(_CurAnswerArrayElementNumber==)
         {
             _StoreAnswerArray[]=_Root->m_Node;
             ++_CurAnswerArrayElementNumber;
         }
         else
         {
             _DISTANCE_TYPE CurNodeDis=_ComputeDistanceFunc(_Value,_Root->m_Node);
             for(unsigned int Idx=_CurAnswerArrayElementNumber-;;--Idx)
             {
                 _DISTANCE_TYPE PreNodeDis=_ComputeDistanceFunc(_Value,_StoreAnswerArray[Idx]);
                 if(_CompareDistanceFunc(CurNodeDis,PreNodeDis))
                 {
                     if(Idx+<_SearchNumber)
                     {
                         _StoreAnswerArray[Idx+]=_StoreAnswerArray[Idx];
                     }
                 }
                 else
                 {
                     if(Idx+<_SearchNumber) _StoreAnswerArray[Idx+]=_Root->m_Node;
                     if(_CurAnswerArrayElementNumber<_SearchNumber)
                     {
                         ++_CurAnswerArrayElementNumber;
                     }
                     break;
                 }
                 if(==Idx)
                 {
                     _StoreAnswerArray[]=_Root->m_Node;
                     if(_CurAnswerArrayElementNumber<_SearchNumber)
                     {
                         ++_CurAnswerArrayElementNumber;
                     }
                     break;
                 }
             }
         }
         if(_Root->m_Left&&_Root->m_Right)
         {
             _DISTANCE_TYPE LSonMinDis=_ComputeMinDistanceFunc(_Value,_Root->m_Left->m_Node);
             _DISTANCE_TYPE RSonMinDis=_ComputeMinDistanceFunc(_Value,_Root->m_Right->m_Node);
             _DISTANCE_TYPE CurMaxDis=_ComputeDistanceFunc(
                 _Value,_StoreAnswerArray[_CurAnswerArrayElementNumber-]);

             if(_CompareDistanceFunc(LSonMinDis,RSonMinDis))
             {
                 if(_CurAnswerArrayElementNumber<_SearchNumber||_CompareDistanceFunc(LSonMinDis,CurMaxDis))
                 {
                     SearchKNear(_Root->m_Left,_Value,_StoreAnswerArray,_CurAnswerArrayElementNumber,
                         _SearchNumber,_CompareDistanceFunc,_ComputeMinDistanceFunc,
                         _ComputeDistanceFunc);
                 }
                 CurMaxDis=_ComputeDistanceFunc(
                     _Value,_StoreAnswerArray[_CurAnswerArrayElementNumber-]);
                 if(_CurAnswerArrayElementNumber<_SearchNumber||_CompareDistanceFunc(RSonMinDis,CurMaxDis))
                 {
                     SearchKNear(_Root->m_Right,_Value,_StoreAnswerArray,_CurAnswerArrayElementNumber,
                         _SearchNumber,_CompareDistanceFunc,_ComputeMinDistanceFunc,
                         _ComputeDistanceFunc);
                 }
             }
             else
             {
                 if(_CurAnswerArrayElementNumber<_SearchNumber||_CompareDistanceFunc(RSonMinDis,CurMaxDis))
                 {
                     SearchKNear(_Root->m_Right,_Value,_StoreAnswerArray,_CurAnswerArrayElementNumber,
                         _SearchNumber,_CompareDistanceFunc,_ComputeMinDistanceFunc,
                         _ComputeDistanceFunc);
                 }
                 CurMaxDis=_ComputeDistanceFunc(
                     _Value,_StoreAnswerArray[_CurAnswerArrayElementNumber-]);
                 if(_CurAnswerArrayElementNumber<_SearchNumber||_CompareDistanceFunc(LSonMinDis,CurMaxDis))
                 {
                     SearchKNear(_Root->m_Left,_Value,_StoreAnswerArray,_CurAnswerArrayElementNumber,
                         _SearchNumber,_CompareDistanceFunc,_ComputeMinDistanceFunc,
                         _ComputeDistanceFunc);
                 }
             }
         }
         else if(_Root->m_Left)
         {
             _DISTANCE_TYPE LSonMinDis=_ComputeMinDistanceFunc(_Value,_Root->m_Left->m_Node);
             _DISTANCE_TYPE CurMaxDis=_ComputeDistanceFunc(
                 _Value,_StoreAnswerArray[_CurAnswerArrayElementNumber-]);
             if(_CurAnswerArrayElementNumber<_SearchNumber||_CompareDistanceFunc(LSonMinDis,CurMaxDis))
             {
                 SearchKNear(_Root->m_Left,_Value,_StoreAnswerArray,_CurAnswerArrayElementNumber,
                     _SearchNumber,_CompareDistanceFunc,_ComputeMinDistanceFunc,
                     _ComputeDistanceFunc);
             }
         }
         else if(_Root->m_Right)
         {
             _DISTANCE_TYPE RSonMinDis=_ComputeMinDistanceFunc(_Value,_Root->m_Right->m_Node);
             _DISTANCE_TYPE CurMaxDis=_ComputeDistanceFunc(
                 _Value,_StoreAnswerArray[_CurAnswerArrayElementNumber-]);
             if(_CurAnswerArrayElementNumber<_SearchNumber||_CompareDistanceFunc(RSonMinDis,CurMaxDis))
             {
                 SearchKNear(_Root->m_Right,_Value,_StoreAnswerArray,_CurAnswerArrayElementNumber,
                     _SearchNumber,_CompareDistanceFunc,_ComputeMinDistanceFunc,
                     _ComputeDistanceFunc);
             }
         }
     }

     unsigned int StoreSubtreeNodeIntoArray(TreeNode* _Root,_NodeType* _NodeArray)
     {
         if(_Root==nullptr) return ;
         unsigned int NodeArraySize=;
         std::queue<TreeNode*> TmpQue;
         TmpQue.push(_Root);
         _NodeArray[NodeArraySize++]=_Root->m_Node;
         while(!TmpQue.empty())
         {
             TreeNode* Tmp=TmpQue.front(); TmpQue.pop();
             if(Tmp->m_Left)
             {
                 TmpQue.push(Tmp->m_Left);
                 _NodeArray[NodeArraySize++]=Tmp->m_Left->m_Node;
             }
             if(Tmp->m_Right)
             {
                 TmpQue.push(Tmp->m_Right);
                 _NodeArray[NodeArraySize++]=Tmp->m_Right->m_Node;
             }
         }
         return NodeArraySize;
     }

     void ClearSubTrees(TreeNode* _Root)
     {
         if(_Root)
         {
             if(_Root->m_Left) ClearSubTrees(_Root->m_Left);
             if(_Root->m_Right) ClearSubTrees(_Root->m_Right);
             delete _Root;
         }
     }

     TreeNode* RebuildTree(TreeNode* _Root,const unsigned int _Dimension)
     {
         _NodeType* TmpPool=new _NodeType[_Root->m_Size];
         unsigned int TmpPoolSize=StoreSubtreeNodeIntoArray(_Root,TmpPool);
         ClearSubTrees(_Root);
         if(TmpPoolSize==) return nullptr;
         TreeNode* Tmp=BuildGroup(TmpPool,,TmpPoolSize-,_Dimension);
         delete[] TmpPool;
         return Tmp;
     }

     template<class _NodeTypeBegin>
     TreeNode* BuildGroup(_NodeTypeBegin _TmpPool,const unsigned int _Left,const unsigned int _Right,const unsigned int _CurLayerDimention)
        {
         if(_Left>_Right) return nullptr;

         const unsigned int MidPos=(_Left+_Right)>>;
         std::nth_element(_TmpPool+_Left,_TmpPool+MidPos,_TmpPool+_Right+,
             m_CompareFuncs[_CurLayerDimention]);

         TreeNode* CurNode=NewNode(_TmpPool[MidPos]);
         if(_Left+<=MidPos)
         {
             CurNode->m_Left=BuildGroup(_TmpPool,_Left,MidPos-,(_CurLayerDimention+)%_DIMENSION);
         }
         CurNode->m_Right=BuildGroup(_TmpPool,MidPos+,_Right,(_CurLayerDimention+)%_DIMENSION);

         PushUp(CurNode);
         return CurNode;
     }

     void PushUp(TreeNode* _Root)
     {
         if(_Root)
         {
             _Root->m_Size=+SonSize(_Root->m_Left)+SonSize(_Root->m_Right);
             _NodeType *LsonNode=_Root->m_Left?&(_Root->m_Left->m_Node):nullptr;
             _NodeType *RsonNode=_Root->m_Right?&_Root->m_Right->m_Node:nullptr;
             m_UpdateFunc(&_Root->m_Node,LsonNode,RsonNode);
         }
     }

     TreeNode* Insert(
         TreeNode* _Root,
         TreeNode* _InsertNode,
         const unsigned int _CurLayerDimention,
         TreeNode* &_BadTreeNode,
         TreeNode* &_BadTreeNodeParent,
         unsigned int& _BadTreeNodeDimension)
     {
         if(nullptr==_Root)
         {
             PushUp(_InsertNode);
             return _InsertNode;
         }
         if(m_CompareFuncs[_CurLayerDimention](_InsertNode->m_Node,_Root->m_Node))
         {
             _Root->m_Left=Insert(_Root->m_Left,_InsertNode,(_CurLayerDimention+)%_DIMENSION,
                 _BadTreeNode,_BadTreeNodeParent,_BadTreeNodeDimension);
         }
         else
         {
             _Root->m_Right=Insert(_Root->m_Right,_InsertNode,(_CurLayerDimention+)%_DIMENSION,
                 _BadTreeNode,_BadTreeNodeParent,_BadTreeNodeDimension);
         }

         PushUp(_Root);

         if(_BadTreeNode==nullptr)
         {
             if(IsSeriousBadTree(_Root))
             {
                 _BadTreeNode=_Root;
                 _BadTreeNodeDimension=_CurLayerDimention;
             }
         }
         else if(_BadTreeNode==_Root->m_Left||_BadTreeNode==_Root->m_Right)
         {
             _BadTreeNodeParent=_Root;
         }
         return _Root;
     }

     unsigned int SonSize(TreeNode* _Node)
     {
         if(_Node==nullptr) return ;
         return _Node->m_Size;
     }

     bool IsSeriousBadTree(TreeNode* _Root)
     {
         if(SonSize(_Root)==) return false;
         return std::max(SonSize(_Root->m_Left),SonSize(_Root->m_Right))
             >=(unsigned int)(SonSize(_Root)*_REBUILD_ALPHA/)+;
     }

     TreeNode* m_Root;
     _CompareFuncType* m_CompareFuncs[_DIMENSION];
     _UpdateFuncType* m_UpdateFunc;
 };

下面是一个简单的测试代码，两个点之间的距离定义为曼哈顿距离。

 struct node
 {
     int x,y;
     int MinX,MaxX,MinY,MaxY;

     node(int _x=,int _y=):x(_x),y(_y) {}

 };

 typedef int Func(const node&,const node&);
 typedef void Func1(node*,node*,node*);

 int cmp0(const node &a,const node &b)
 {
     return a.x<b.x;
 }

 int cmp1(const node &a,const node &b)
 {
     return a.y<b.y;
 }

 void pushUp(node *Fa,node *lson,node *rson)
 {
     Fa->MinX=Fa->MaxX=Fa->x;
     Fa->MinY=Fa->MaxY=Fa->y;
     for(int i=;i<;++i)
     {
         node* p=i==?lson:rson;
         if(!p) continue;

         Fa->MinX=min(Fa->MinX,p->MinX);
         Fa->MaxX=max(Fa->MaxX,p->MaxX);
         Fa->MinY=min(Fa->MinY,p->MinY);
         Fa->MaxY=max(Fa->MaxY,p->MaxY);
     }
 }

 int caldis(const node &a,const node &b)
 {
     return abs(a.x-b.x)+abs(a.y-b.y);
 }

 int calMinDis(const node &a,const node &b)
 {
     int xx=;
     int yy=;
     if(a.x<b.MinX) xx=b.MinX-a.x;
     else if(a.x>b.MaxX) xx=a.x-b.MaxX;

     if(a.y<b.MinY) yy=b.MinY-a.y;
     else if(a.y>b.MaxY) yy=a.y-b.MaxY;

     return xx+yy;
 }

 int cmp(int x,int y)
 {
     return x<y;
 }

 int main()
 {
     Func* funs[]={cmp0,cmp1};
     KDTree<node,Func,Func1,,int> *T=
         new KDTree<node,Func,Func1,,int>(funs,pushUp);

     node a[]={node(,),node(-,),node(,),node(,)};
     T->insert(a,);
     T->insert(node(,));
     unsigned int Num=T->searchKNear(node(,),a,,cmp,calMinDis,caldis);
     for(unsigned int Idx=;Idx<Num;++Idx)
     {
         printf("%d %d\n",a[Idx].x,a[Idx].y);
     }
     /**
     5 6
     5 9
     2 3
     **/
 }