kdTree.h

   1 #ifndef KDTREE_HEADER_
   2 #define KDTREE_HEADER_
   3
   4 #ifdef NDEBUG
   5         #include <cstring> // memcpy
   6 #endif
   7
   8 namespace NOSPACE {
   9         using namespace std;
  10 }
  11
  12 /** Routines for computing with T[length] number fields */
  13 namespace FieldMath {
  14         /** Calls transformer(x,y) for every x from [i1;iEnd1) and corresponding y from [i2;?) */
  15         template<class I1,class I2,class Transf> inline
  16         Transf transform2( I1 i1, I1 iEnd1, I2 i2, Transf transformer ) {
  17                 for (; i1!=iEnd1; ++i1,++i2)
  18                         transformer(*i1,*i2);
  19                 return transformer;
  20         }
  21         /** Analogous to ::transform2 */
  22         template<class I1,class I2,class I3,class Transf> inline
  23         Transf transform3( I1 i1, I1 iEnd1, I2 i2, I3 i3, Transf transformer ) {
  24                 for (; i1!=iEnd1; ++i1,++i2,++i3)
  25                         transformer(*i1,*i2,*i3);
  26                 return transformer;
  27         }
  28
  29         /** Means b[i]=a[i]; Only meant for POD types */
  30         template<class T> inline T* assign(const T *a,int length,T *b) {
  31         #ifndef NDEBUG
  32                 copy( a, a+length, b );                         // debugging version uses STL's copy
  33         #else
  34                 memcpy( b, a, length*sizeof(T) );       // release version uses C's memory-copy (faster)
  35         #endif
  36                 return b;
  37         }
  38
  39         namespace NOSPACE {
  40                 /** A helper transforming structure, only to be used in ::moveToBounds_copy */
  41                 template<class T,bool CheckNaNs> struct MoveToBounds {
  42                         T sqrError; ///< square error accumulated so far
  43
  44                         /** Only sets ::sqrError to zero */
  45                         MoveToBounds()
  46                         : sqrError(0) {}
  47
  48                         /** Moves one coordinate of a point (in \p point) to a bounding box
  49                          *      (in \p bounds) accumulating ::sqrError and storing result (in \p result) */
  50                         void operator()(const T &point,const T bounds[2],T &result) {
  51                                 if ( CheckNaNs && isNaN(point) )
  52                                         return;
  53                                 if ( point < bounds[0] ) {
  54                                         sqrError+= sqr(bounds[0]-point);
  55                                         result= bounds[0];
  56                                 } else
  57                                 if ( point > bounds[1] ) {
  58                                         sqrError+= sqr(point-bounds[1]);
  59                                         result= bounds[1];
  60                                 } else
  61                                         result= point;
  62                         }
  63                 };
  64         }
  65         /** Copy_moves a vector (\p point) to the nearest point (\p result)
  66          *      within bounds (\p bounds) and returns SE (distance^2) */
  67         template<class T,bool CheckNaNs> inline
  68         T moveToBounds_copy(const T *point,const T (*bounds)[2],int length,T *result) {
  69                 return transform3( point, point+length, bounds, result
  70                         , MoveToBounds<T,CheckNaNs>() ) .sqrError;
  71         }
  72 }
  73
  74 template<class T> class KDBuilder;
  75 /** A generic static KD-tree \todo */
  76 template<class T> class KDTree {
  77 public:
  78         friend class KDBuilder<T>;
  79         typedef KDBuilder<T> Builder;
  80 protected:
  81         /** Represents one node of the KD-tree */
  82         struct Node {
  83                 int coord;      ///< The coordinate along which the tree splits in this node
  84                 T threshold;/**< The threshold of the split
  85                                          *       (left[::coord] <= ::threshold <= right[::coord]) */
  86         };
  87         typedef T (*Bounds)[2];
  88
  89 public:
  90         const int depth ///  The depth of the tree = ::log2ceil(::count)
  91         , length                ///      The length of the vectors
  92         , count;                ///< The number of the vectors
  93 protected:
  94         Node *nodes;    ///< The array of the tree-nodes (heap-like topology of the tree)
  95         int *dataIDs;   ///< Data IDs for children of "leaf" nodes
  96         Bounds bounds;  ///< The bounding box for all the data
  97
  98         /** Prepares to build a new KD-tree from \p count_ vectors of \p length_ elements */
  99         KDTree(int length_,int count_)
 100         : depth( log2ceil(count_) ), length(length_), count(count_)
 101         , nodes( new Node[count_] ), dataIDs( new int[count_] ), bounds( new T[length_][2] ) {
 102         }
 103
 104         /** Copy constructor with moving semantics (destroys its argument) */
 105         KDTree(KDTree &other)
 106         : depth(other.depth), length(other.length), count(other.count)
 107         , nodes(other.nodes), dataIDs(other.dataIDs), bounds(other.bounds) {
 108                 other.nodes= 0;
 109                 other.dataIDs= 0;
 110                 other.bounds= 0;
 111         }
 112
 113         /** Takes an index of a "leaf" node (past the end of ::nodes)
 114          *      and returns the appropriate data ID */
 115         int leafID2dataID(int leafID) const {
 116                 ASSERT( count<=leafID && leafID<2*count );
 117                 int index= leafID-powers[depth];
 118                 if (index<0)
 119                         index+= count; // it is on the shallower side of the tree
 120                 ASSERT( 0<=index && index<count );
 121                 return dataIDs[index];
 122         }
 123
 124 public:
 125         /** Only frees the memory */
 126         ~KDTree() {
 127         //      clean up
 128                 delete[] nodes;
 129                 delete[] dataIDs;
 130                 delete[] bounds;
 131         }
 132
 133         /** Manages a heap from nodes of a KDTree.
 134          *      It returns vectors (their indices) in the order of ascending distance (SE)
 135          *      from a given fixed point. It can compute a lower bound of the SEs of the remaining
 136          *      vectors at any time. */
 137         class PointHeap {
 138                 /** One element of the ::heap representing a node in the KDTree ::kd */
 139                 struct HeapNode {
 140                         int nodeIndex;  ///< Index of the node in ::kd
 141                         T *data;                /**< Stores the SE and the coordinates of the nearest point
 142                                                          *       (to this node's bounding box) */
 143                         /** No-init constructor */
 144                         HeapNode() {}
 145                         /** Initializes the members from the parameters */
 146                         HeapNode(int nodeIndex_,T *data_): nodeIndex(nodeIndex_), data(data_) {}
 147
 148                         /** Returns reference to the SE of the nearest point to this node's bounding box */
 149                         T& getSE()                      { return *data; }
 150                         /** Returns the SE of the nearest point to this node's bounding box */
 151                         T getSE() const         { return *data; }
 152                         /** Returns pointer to the nearest point to this node's bounding box */
 153                         T* getNearest()         { return data+1; }
 154                 };
 155                 /** Defines the order of ::heap - ascending according to ::getSE */
 156                 struct HeapOrder {
 157                         bool operator()(const HeapNode &a,const HeapNode &b)
 158                                 { return a.getSE() > b.getSE(); }
 159                 };
 160
 161                 const KDTree &kd;                       ///< Reference to the KDTree we operate on
 162                 const T* const point;           ///< Pointer to the point we are trying to approach
 163                 vector<HeapNode> heap;          ///< The current heap of the tree nodes
 164                 BulkAllocator<T> allocator;     ///< The allocator for HeapNode::data
 165         public:
 166                 /** Builds the heap from a KDTree \p tree and vector \p point_
 167                  *      (they've got to remain valid until the destruction of this instance) */
 168                 PointHeap(const KDTree &tree,const T *point_,bool checkNaNs)
 169                 : kd(tree), point(point_) {
 170                         ASSERT(point);
 171                 //      create the root heap-node
 172                         HeapNode rootNode( 1, allocator.makeField(kd.length+1) );
 173                 //      compute the nearest point within the global bounds and corresponding SE
 174                         using namespace FieldMath;
 175                         rootNode.getSE()= checkNaNs
 176                                 ? moveToBounds_copy<T,true> ( point, kd.bounds, kd.length, rootNode.getNearest() )
 177                                 : moveToBounds_copy<T,false>( point, kd.bounds, kd.length, rootNode.getNearest() );
 178                 //      push it onto the heap (and reserve more to speed up the first leaf-gettings)
 179                         heap.reserve(kd.depth*2);
 180                         heap.push_back(rootNode);
 181                 }
 182
 183                 /** Returns whether the heap is empty ( !isEmpty() is needed for all other methods) */
 184                 bool isEmpty()
 185                         { return heap.empty(); }
 186
 187                 /** Returns the SE of the top node (always equals the SE of the next leaf) */
 188                 T getTopSE() {
 189                         ASSERT( !isEmpty() );
 190                         return heap[0].getSE();
 191                 }
 192
 193                 /** Removes a leaf, returns the matching vector's index,
 194                  *      assumes it's safe to discard nodes further than \p maxSE */
 195                 template<bool CheckNaNs> int popLeaf(T maxSE) {
 196                 //      ensure a leaf is on the top of the heap and get its ID
 197                         makeTopLeaf<CheckNaNs>(maxSE);
 198                         int result= kd.leafID2dataID( heap.front().nodeIndex );
 199                 //      remove the top from the heap (no need to free the memory - ::allocator)
 200                         pop_heap( heap.begin(), heap.end(), HeapOrder() );
 201                         heap.pop_back();
 202                         return result;
 203                 }
 204         protected:
 205                 /** Divides the top nodes until there's a leaf on the top
 206                  *      assumes it's safe to discard nodes further than \p maxSE */
 207                 template<bool CheckNaNs> void makeTopLeaf(T maxSE);
 208
 209         }; // PointHeap class
 210 }; // KDTree class
 211
 212
 213 template<class T> template<bool CheckNaNs>
 214 void KDTree<T>::PointHeap::makeTopLeaf(T maxSE) {
 215         ASSERT( !isEmpty() );
 216 //      exit if there's a leaf on the top already
 217         if ( heap[0].nodeIndex >= kd.count )
 218                 return;
 219         PtrInt oldHeapSize= heap.size();
 220         HeapNode heapRoot= heap[0]; // making a local working copy of the top of the heap
 221
 222         do { // while heapRoot isn't leaf ... while ( heapRoot.nodeIndex<kd.count )
 223                 const Node &node= kd.nodes[heapRoot.nodeIndex];
 224         //      now replace the node with its two children:
 225         //              one of them will have the same SE (replaces its parent on the top),
 226         //              the other one can have higher SE (push_back-ed on the heap)
 227
 228                 bool validCoord= !CheckNaNs || !isNaN(point[node.coord]);
 229         //      the higher SE can be computed from the parent heap-node
 230                 T newSE;
 231                 bool goRight; // will the heapRoot represent the right child or the left one?
 232                 if (validCoord) {
 233                         Real oldDiff= Real(point[node.coord]) - heapRoot.getNearest()[node.coord];
 234                         Real newDiff= Real(point[node.coord]) - node.threshold;
 235                         goRight= newDiff>0;
 236                         newSE= heapRoot.getSE() - sqr(oldDiff) + sqr(newDiff);
 237                         ASSERT( newSE >= heapRoot.getSE() );
 238                 } else {
 239                         newSE= heapRoot.getSE();
 240                         goRight= false; // both boolean values are possible
 241                 }
 242         //      the root will represent its closer child - neither nearest point nor SE changes
 243                 heapRoot.nodeIndex= heapRoot.nodeIndex*2 + goRight;
 244         //      if the new SE is too high, continue (omit the child with this big SE)
 245                 if (newSE>maxSE)
 246                         continue;
 247         //      create a new heap-node, allocate it's data and assign index of the other child
 248                 HeapNode newHNode;
 249                 newHNode.data= allocator.makeField(kd.length+1);
 250                 newHNode.getSE()= newSE;
 251                 newHNode.nodeIndex= heapRoot.nodeIndex-goRight+!goRight;
 252         //      the nearest point of the new heap-node only differs in one coordinate
 253                 FieldMath::assign( heapRoot.getNearest(), kd.length, newHNode.getNearest() );
 254                 if (validCoord)
 255                         newHNode.getNearest()[node.coord]= node.threshold;
 256         //      add the new node to the back, restore the heap-property later
 257                 heap.push_back(newHNode);
 258
 259         } while ( heapRoot.nodeIndex < kd.count );
 260
 261         heap[0]= heapRoot; // restoring the working copy of the heap's top node
 262 //      restore the heap-property on the added nodes
 263         typename vector<HeapNode>::iterator it= heap.begin()+oldHeapSize;
 264         do
 265                 push_heap( heap.begin(), it, HeapOrder() );
 266         while ( it++ != heap.end() );
 267 } // KDTree<T>::PointHeap::makeTopLeaf<CheckNaNs>() method
 268
 269 /** Derived type used to construct KDTree instances (::makeTree static method) */
 270 template<class T> class KDBuilder: public KDTree<T> {
 271 public:
 272         typedef KDTree<T> Tree;
 273         typedef T BoundsPair[2];
 274         typedef typename Tree::Bounds Bounds;
 275         /** Type for a method that chooses which coordinate to split */
 276         typedef int (KDBuilder::*CoordChooser)
 277                 (int nodeIndex,int *beginIDs,int *endIDs,int depthLeft) const;
 278 protected:
 279         using Tree::depth;      using Tree::length;             using Tree::count;
 280         using Tree::nodes;      using Tree::dataIDs;    using Tree::bounds;
 281
 282         const T *data;
 283         const CoordChooser chooser;
 284         mutable Bounds chooserTmp;
 285
 286         KDBuilder(const T *data_,int length,int count,CoordChooser chooser_)
 287         : Tree(length,count), data(data_), chooser(chooser_), chooserTmp(0) {
 288                 ASSERT( length>0 && count>0 && chooser && data );
 289         //      create the index-vector, coumpute the bounding box, build the tree
 290                 for (int i=0; i<count; ++i)
 291                         dataIDs[i]= i;
 292                 getBounds(bounds);
 293                 if (count>1)
 294                         buildNode(1,dataIDs,dataIDs+count,depth);
 295                 delete[] chooserTmp;
 296                 DEBUG_ONLY( chooserTmp= 0; data= 0; )
 297         }
 298
 299         /** Creates bounds containing one value */
 300         struct NewBounds {
 301                 void operator()(const T &val,BoundsPair &bounds) const
 302                         { bounds[0]= bounds[1]= val; }
 303         };
 304         /** Expands a valid bounding box to contain a point (one coordinate at once) */
 305         struct BoundsExpander {
 306                 void operator()(const T &val,BoundsPair &bounds) const {
 307                         if ( val < bounds[0] ) // lower bound
 308                                 bounds[0]= val; else
 309                         if ( val > bounds[1] ) // upper bound
 310                                 bounds[1]= val;
 311                 }
 312         };
 313         /** Computes the bounding box of ::count vectors with length ::length stored in ::data.
 314          *      The vectors in ::data are stored linearly, \p boundsRes should be preallocated */
 315         void getBounds(Bounds boundsRes) const {
 316                 using namespace FieldMath;
 317                 ASSERT(length>0);
 318         //      make the initial bounds only contain the first point
 319                 transform2(data,data+length,boundsRes,NewBounds());
 320                 int count= Tree::count;
 321                 const T *nowData= data;
 322         //      expand the bounds by every point (except for the first one)
 323                 while (--count) {
 324                         nowData+= length;
 325                         transform2( nowData, nowData+length, boundsRes, BoundsExpander() );
 326                 }
 327         }
 328         /** Like ::getBounds, but it only works on vectors with indices from [\p beginIDs;\p endIDs)
 329          *      instead of [0;::count), \p boundsRes should be preallocated to store the result */
 330         void getBounds(const int *beginIDs,const int *endIDs,Bounds boundsRes) const {
 331                 using namespace FieldMath;
 332                 ASSERT(endIDs>beginIDs);
 333         //      make the initial bounds only contain the first point
 334                 const T *begin= data + *beginIDs*length;
 335                 transform2(begin,begin+length,boundsRes,NewBounds());
 336         //      expand the bounds by every point (except for the first one)
 337                 while ( ++beginIDs != endIDs ) {
 338                         begin= data + *beginIDs*length;
 339                         transform2( begin, begin+length, boundsRes, BoundsExpander() );
 340                 }
 341         }
 342
 343         /** Recursively builds node \p nodeIndex and its subtree of depth \p depthLeft
 344         *       (including leaves), operates on data \p data in the range [\p beginIDs,\p endIDs) */
 345         void buildNode(int nodeIndex,int *beginIDs,int *endIDs,int depthLeft);
 346
 347 public:
 348         /** Builds a KDTree from \p count vectors of length \p length stored in \p data,
 349          *      splitting the nodes by \p chooser CoordChooser */
 350         static Tree* makeTree(const T *data,int length,int count,CoordChooser chooser) {
 351                 KDBuilder builder(data,length,count,chooser);
 352         //      moving only the necesarry data (pointers) into a new copy
 353                 return new Tree(builder);
 354         }
 355
 356         /** CoordChooser choosing the longest coordinate of the bounding box of the current interval */
 357         int choosePrecise(int nodeIndex,int *beginIDs,int *endIDs,int /*depthLeft*/) const;
 358         /** CoordChooser choosing the coordinate only according to the depth */
 359         int chooseFast(int /*nodeIndex*/,int* /*beginIDs*/,int* /*endIDs*/,int depthLeft) const
 360                 { return depthLeft%length; }
 361         /** CoordChooser choosing a random coordinate */
 362         int chooseRand(int /*nodeIndex*/,int* /*beginIDs*/,int* /*endIDs*/,int /*depthLeft*/) const
 363                 { return rand()%length; }
 364         /** CoordChooser - like ::choosePrecise, but doesn't compute the real bounding box,
 365          *      only approximates it by splitting the parent's box (a little less accurate, but much faster) */
 366         int chooseApprox(int nodeIndex,int* /*beginIDs*/,int* /*endIDs*/,int depthLeft) const;
 367 }; // KDBuilder class
 368
 369
 370
 371 namespace NOSPACE {
 372         /** Finds the longest coordinate of a bounding box,     only to be used in for_each calls */
 373         template<class T> struct MaxDiffCoord {
 374                 typedef typename KDBuilder<T>::BoundsPair BoundsPair;
 375
 376                 T maxDiff;              ///< the maximal difference value
 377                 int bestIndex   ///  the index where ::maxDiff occured
 378                 , nextIndex;    ///< next index to be checked
 379
 380                 /** Initializes from the 0-th index value */
 381                 MaxDiffCoord(const BoundsPair& bounds0)
 382                 : maxDiff(bounds0[1]-bounds0[0]), bestIndex(0), nextIndex(1) {}
 383
 384                 /** To be called successively for indices from 1-st on */
 385                 void operator()(const BoundsPair& bounds_i) {
 386                         T diff= bounds_i[1]-bounds_i[0];
 387                         if (diff>maxDiff) {
 388                                 bestIndex= nextIndex;
 389                                 maxDiff= diff;
 390                         }
 391                         ++nextIndex;
 392                 }
 393         };
 394 }
 395 template<class T> int KDBuilder<T>
 396 ::choosePrecise(int nodeIndex,int *beginIDs,int *endIDs,int) const {
 397         ASSERT( nodeIndex>0 && beginIDs && endIDs && beginIDs<endIDs );
 398 //      temporary storage for computed bounding box
 399         BoundsPair boundsStorage[length];
 400         const BoundsPair *localBounds;
 401 //      find out the bounding box
 402         if ( nodeIndex>1 ) { // compute the bounding box
 403                 localBounds= boundsStorage;
 404                 getBounds( beginIDs, endIDs, boundsStorage );
 405         } else // we are in the root -> we can use already computed bounds
 406                 localBounds= this->bounds;
 407 //      find and return the longest coordinate
 408         MaxDiffCoord<T> mdc= for_each
 409                 ( localBounds+1, localBounds+length, MaxDiffCoord<T>(localBounds[0]) );
 410         ASSERT( mdc.nextIndex == length );
 411         return mdc.bestIndex;
 412 }
 413
 414 template<class T> int KDBuilder<T>::chooseApprox(int nodeIndex,int*,int*,int) const {
 415         using namespace FieldMath;
 416         ASSERT(nodeIndex>0);
 417
 418         int myDepth= log2ceil(nodeIndex+1)-1;
 419         if (!myDepth) { // I'm in the root - copy the bounds
 420                 ASSERT(nodeIndex==1);
 421                 chooserTmp= new BoundsPair[length*(depth+1)]; // allocate my temporary bound-array
 422                 assign( bounds, length, chooserTmp );
 423         }
 424
 425         Bounds myBounds= chooserTmp+length*myDepth;
 426         if (myDepth) { // I'm not the root - copy parent's bounds and modify them
 427                 const typename Tree::Node &parent= nodes[nodeIndex/2];
 428                 Bounds parentBounds= myBounds-length;
 429                 if (nodeIndex%2) {      // I'm the right son -> bounds on this level not initialized
 430                         assign( parentBounds, length, myBounds );               // copying parent bounds
 431                         myBounds[parent.coord][0]= parent.threshold;    // adjusting the lower bound
 432                 } else // I'm the left son
 433                         if ( nodeIndex+1 < count ) { // almost the same as brother -> only adjust the coordinate
 434                                 myBounds[parent.coord][0]= parentBounds[parent.coord][0];
 435                                 myBounds[parent.coord][1]= parent.threshold;
 436                         } else { // I've got no brother
 437                                 ASSERT( nodeIndex+1 == count );
 438                                 assign( parentBounds, length, myBounds );
 439                                 myBounds[parent.coord][1]= parent.threshold;
 440                         }
 441         }
 442 //      find out the widest dimension
 443         MaxDiffCoord<T> mdc= for_each( myBounds+1, myBounds+length, MaxDiffCoord<T>(myBounds[0]) );
 444         ASSERT( mdc.nextIndex == length );
 445         return mdc.bestIndex;
 446 }
 447
 448 namespace NOSPACE {
 449         /** Compares vectors (given by their indices) according to a given coordinate */
 450         template<class T> class IndexComparator {
 451                 const T *data;  ///< pointer to the right coordinate of the first vector (of index 0)
 452                 int length;             ///< length of the vectors in ::data
 453         public:
 454                 IndexComparator(const T *data_,int length_,int index_)
 455                 : data(data_+index_), length(length_) {}
 456                 bool operator()(int a,int b) const
 457                         { return data[a*length] < data[b*length]; }
 458         };
 459 }
 460 template<class T> void KDBuilder<T>
 461 ::buildNode(int nodeIndex,int *beginIDs,int *endIDs,int depthLeft) {
 462         int count= endIDs-beginIDs; // owershadowing Tree::count
 463 //      check we've got at least one vector and the depth&count are adequate to each other
 464         ASSERT( count>=2 && powers[depthLeft-1]<count && count<=powers[depthLeft] );
 465         --depthLeft;
 466 //      find out where to split - how many items should be on the left to have the heap-shape
 467         bool shallowRight= ( count <= powers[depthLeft]+powers[depthLeft-1] );
 468         int *middle= shallowRight
 469                 ? endIDs-powers[depthLeft-1]
 470                 : beginIDs+powers[depthLeft];
 471 //      find out the dividing coordinate and find the "median" in this coordinate
 472         int coord= (this->*chooser)(nodeIndex,beginIDs,endIDs,depthLeft);
 473         nth_element( beginIDs, middle , endIDs, IndexComparator<T>(data,length,coord) );
 474 //      fill the node's data (dividing coordinate and its threshold)
 475         nodes[nodeIndex].coord= coord;
 476         nodes[nodeIndex].threshold= data[*middle*length+coord]; // min. value of the right son
 477 //      recurse on both halves (if needed; fall-through switch)
 478         switch (count) {
 479         default: //     we've got enough nodes - build both subtrees (fall through)
 480         //      build the right subtree
 481                 buildNode( 2*nodeIndex+1, middle, endIDs, depthLeft-shallowRight );
 482         case 3: // only a pair in the first half
 483         //      build the left subtree
 484                 buildNode( 2*nodeIndex, beginIDs, middle, depthLeft );
 485         case 2: // nothing needs to be sorted
 486                 ;
 487         }
 488 }
 489
 490 #endif // KDTREE_HEADER_