kdTree.h

Go to the documentation of this file.

#ifndef KDTREE_HEADER_
#define KDTREE_HEADER_

#ifdef NDEBUG
    #include <cstring> // memcpy
#endif

namespace NOSPACE {
    using namespace std;
}

namespace FieldMath {
    template<class I1,class I2,class Transf> inline
    Transf transform2( I1 i1, I1 iEnd1, I2 i2, Transf transformer ) {
        for (; i1!=iEnd1; ++i1,++i2)
            transformer(*i1,*i2);
        return transformer;
    }
    template<class I1,class I2,class I3,class Transf> inline
    Transf transform3( I1 i1, I1 iEnd1, I2 i2, I3 i3, Transf transformer ) {
        for (; i1!=iEnd1; ++i1,++i2,++i3)
            transformer(*i1,*i2,*i3);
        return transformer;
    }

    template<class T> inline T* assign(const T *a,int length,T *b) {
    #ifndef NDEBUG
        copy( a, a+length, b );             // debugging version uses STL's copy
    #else
        memcpy( b, a, length*sizeof(T) );   // release version uses C's memory-copy (faster)
    #endif
        return b;
    }

    namespace NOSPACE {
        template<class T,bool CheckNaNs> struct MoveToBounds {
            T sqrError; 

            MoveToBounds()
            : sqrError(0) {}

            void operator()(const T &point,const T bounds[2],T &result) {
                if ( CheckNaNs && isNaN(point) )
                    return;
                if ( point < bounds[0] ) {
                    sqrError+= sqr(bounds[0]-point);
                    result= bounds[0];
                } else
                if ( point > bounds[1] ) {
                    sqrError+= sqr(point-bounds[1]);
                    result= bounds[1];
                } else
                    result= point;
            }
        };
    }
    template<class T,bool CheckNaNs> inline
    T moveToBounds_copy(const T *point,const T (*bounds)[2],int length,T *result) {
        return transform3( point, point+length, bounds, result
            , MoveToBounds<T,CheckNaNs>() ) .sqrError;
    }
}

template<class T> class KDBuilder;
template<class T> class KDTree {
public:
    friend class KDBuilder<T>;
    typedef KDBuilder<T> Builder;
protected:
    struct Node {
        int coord;  
        T threshold;
    };
    typedef T (*Bounds)[2];

public:
    const int depth 
    , length        
    , count;        
protected:
    Node *nodes;    
    int *dataIDs;   
    Bounds bounds;  

    KDTree(int length_,int count_)
    : depth( log2ceil(count_) ), length(length_), count(count_)
    , nodes( new Node[count_] ), dataIDs( new int[count_] ), bounds( new T[length_][2] ) {
    }

    KDTree(KDTree &other)
    : depth(other.depth), length(other.length), count(other.count)
    , nodes(other.nodes), dataIDs(other.dataIDs), bounds(other.bounds) {
        other.nodes= 0;
        other.dataIDs= 0;
        other.bounds= 0;
    }
    
    int leafID2dataID(int leafID) const {
        ASSERT( count<=leafID && leafID<2*count );
        int index= leafID-powers[depth];
        if (index<0)
            index+= count; // it is on the shallower side of the tree
        ASSERT( 0<=index && index<count );
        return dataIDs[index];
    }

public:
    ~KDTree() {
    //  clean up
        delete[] nodes;
        delete[] dataIDs;
        delete[] bounds;
    }

    class PointHeap {
        struct HeapNode {
            int nodeIndex;  
            T *data;        
            HeapNode() {}
            HeapNode(int nodeIndex_,T *data_): nodeIndex(nodeIndex_), data(data_) {}

            T& getSE()          { return *data; }
            T getSE() const     { return *data; }
            T* getNearest()     { return data+1; }
        };
        struct HeapOrder {
            bool operator()(const HeapNode &a,const HeapNode &b)
                { return a.getSE() > b.getSE(); }
        };

        const KDTree &kd;           
        const T* const point;       
        vector<HeapNode> heap;      
        BulkAllocator<T> allocator; 
    public:
        PointHeap(const KDTree &tree,const T *point_,bool checkNaNs)
        : kd(tree), point(point_) {
            ASSERT(point);
        //  create the root heap-node
            HeapNode rootNode( 1, allocator.makeField(kd.length+1) );
        //  compute the nearest point within the global bounds and corresponding SE
            using namespace FieldMath;
            rootNode.getSE()= checkNaNs
                ? moveToBounds_copy<T,true> ( point, kd.bounds, kd.length, rootNode.getNearest() )
                : moveToBounds_copy<T,false>( point, kd.bounds, kd.length, rootNode.getNearest() );
        //  push it onto the heap (and reserve more to speed up the first leaf-gettings)
            heap.reserve(kd.depth*2);
            heap.push_back(rootNode);
        }

        bool isEmpty()
            { return heap.empty(); }

        T getTopSE() {
            ASSERT( !isEmpty() );
            return heap[0].getSE();
        }

        template<bool CheckNaNs> int popLeaf(T maxSE) {
        //  ensure a leaf is on the top of the heap and get its ID
            makeTopLeaf<CheckNaNs>(maxSE);
            int result= kd.leafID2dataID( heap.front().nodeIndex );
        //  remove the top from the heap (no need to free the memory - ::allocator)
            pop_heap( heap.begin(), heap.end(), HeapOrder() );
            heap.pop_back();
            return result;
        }
    protected:
        template<bool CheckNaNs> void makeTopLeaf(T maxSE);

    }; // PointHeap class
}; // KDTree class


template<class T> template<bool CheckNaNs>
void KDTree<T>::PointHeap::makeTopLeaf(T maxSE) {
    ASSERT( !isEmpty() );
//  exit if there's a leaf on the top already
    if ( heap[0].nodeIndex >= kd.count )
        return;
    PtrInt oldHeapSize= heap.size();
    HeapNode heapRoot= heap[0]; // making a local working copy of the top of the heap

    do { // while heapRoot isn't leaf ... while ( heapRoot.nodeIndex<kd.count )
        const Node &node= kd.nodes[heapRoot.nodeIndex];
    //  now replace the node with its two children:
    //      one of them will have the same SE (replaces its parent on the top),
    //      the other one can have higher SE (push_back-ed on the heap)

        bool validCoord= !CheckNaNs || !isNaN(point[node.coord]);
    //  the higher SE can be computed from the parent heap-node
        T newSE;
        bool goRight; // will the heapRoot represent the right child or the left one?
        if (validCoord) {
            Real oldDiff= Real(point[node.coord]) - heapRoot.getNearest()[node.coord];
            Real newDiff= Real(point[node.coord]) - node.threshold;
            goRight= newDiff>0;
            newSE= heapRoot.getSE() - sqr(oldDiff) + sqr(newDiff);
            ASSERT( newSE >= heapRoot.getSE() );
        } else {
            newSE= heapRoot.getSE();
            goRight= false; // both boolean values are possible 
        }
    //  the root will represent its closer child - neither nearest point nor SE changes
        heapRoot.nodeIndex= heapRoot.nodeIndex*2 + goRight;
    //  if the new SE is too high, continue (omit the child with this big SE)
        if (newSE>maxSE)    
            continue;
    //  create a new heap-node, allocate it's data and assign index of the other child
        HeapNode newHNode;
        newHNode.data= allocator.makeField(kd.length+1);
        newHNode.getSE()= newSE;    
        newHNode.nodeIndex= heapRoot.nodeIndex-goRight+!goRight;
    //  the nearest point of the new heap-node only differs in one coordinate
        FieldMath::assign( heapRoot.getNearest(), kd.length, newHNode.getNearest() );
        if (validCoord)
            newHNode.getNearest()[node.coord]= node.threshold;
    //  add the new node to the back, restore the heap-property later
        heap.push_back(newHNode);
        
    } while ( heapRoot.nodeIndex < kd.count );

    heap[0]= heapRoot; // restoring the working copy of the heap's top node
//  restore the heap-property on the added nodes
    typename vector<HeapNode>::iterator it= heap.begin()+oldHeapSize;
    do
        push_heap( heap.begin(), it, HeapOrder() );
    while ( it++ != heap.end() );
} // KDTree<T>::PointHeap::makeTopLeaf<CheckNaNs>() method

template<class T> class KDBuilder: public KDTree<T> {
public:
    typedef KDTree<T> Tree;
    typedef T BoundsPair[2];
    typedef typename Tree::Bounds Bounds;
    typedef int (KDBuilder::*CoordChooser)
        (int nodeIndex,int *beginIDs,int *endIDs,int depthLeft) const;
protected:
    using Tree::depth;  using Tree::length;     using Tree::count;
    using Tree::nodes;  using Tree::dataIDs;    using Tree::bounds;

    const T *data;
    const CoordChooser chooser;
    mutable Bounds chooserTmp;

    KDBuilder(const T *data_,int length,int count,CoordChooser chooser_)
    : Tree(length,count), data(data_), chooser(chooser_), chooserTmp(0) {
        ASSERT( length>0 && count>0 && chooser && data );
    //  create the index-vector, coumpute the bounding box, build the tree
        for (int i=0; i<count; ++i)
            dataIDs[i]= i;
        getBounds(bounds);
        if (count>1)
            buildNode(1,dataIDs,dataIDs+count,depth);
        delete[] chooserTmp;
        DEBUG_ONLY( chooserTmp= 0; data= 0; )
    }

    struct NewBounds {
        void operator()(const T &val,BoundsPair &bounds) const
            { bounds[0]= bounds[1]= val; }
    };
    struct BoundsExpander {
        void operator()(const T &val,BoundsPair &bounds) const {
            if ( val < bounds[0] ) // lower bound
                bounds[0]= val; else
            if ( val > bounds[1] ) // upper bound
                bounds[1]= val;
        }
    };
    void getBounds(Bounds boundsRes) const {
        using namespace FieldMath;
        ASSERT(length>0);
    //  make the initial bounds only contain the first point
        transform2(data,data+length,boundsRes,NewBounds());
        int count= Tree::count;
        const T *nowData= data;
    //  expand the bounds by every point (except for the first one)
        while (--count) {
            nowData+= length;
            transform2( nowData, nowData+length, boundsRes, BoundsExpander() );
        }
    }
    void getBounds(const int *beginIDs,const int *endIDs,Bounds boundsRes) const {
        using namespace FieldMath;
        ASSERT(endIDs>beginIDs);
    //  make the initial bounds only contain the first point
        const T *begin= data + *beginIDs*length;
        transform2(begin,begin+length,boundsRes,NewBounds());
    //  expand the bounds by every point (except for the first one)
        while ( ++beginIDs != endIDs ) {
            begin= data + *beginIDs*length;
            transform2( begin, begin+length, boundsRes, BoundsExpander() );
        }
    }

    void buildNode(int nodeIndex,int *beginIDs,int *endIDs,int depthLeft);

public:
    static Tree* makeTree(const T *data,int length,int count,CoordChooser chooser) {
        KDBuilder builder(data,length,count,chooser);
    //  moving only the necesarry data (pointers) into a new copy
        return new Tree(builder);
    }
    
    int choosePrecise(int nodeIndex,int *beginIDs,int *endIDs,int /*depthLeft*/) const;
    int chooseFast(int /*nodeIndex*/,int* /*beginIDs*/,int* /*endIDs*/,int depthLeft) const
        { return depthLeft%length; }
    int chooseRand(int /*nodeIndex*/,int* /*beginIDs*/,int* /*endIDs*/,int /*depthLeft*/) const
        { return rand()%length; }
    int chooseApprox(int nodeIndex,int* /*beginIDs*/,int* /*endIDs*/,int depthLeft) const;
}; // KDBuilder class



namespace NOSPACE {
    template<class T> struct MaxDiffCoord {
        typedef typename KDBuilder<T>::BoundsPair BoundsPair;

        T maxDiff;      
        int bestIndex   
        , nextIndex;    

        MaxDiffCoord(const BoundsPair& bounds0)
        : maxDiff(bounds0[1]-bounds0[0]), bestIndex(0), nextIndex(1) {}

        void operator()(const BoundsPair& bounds_i) {
            T diff= bounds_i[1]-bounds_i[0];
            if (diff>maxDiff) {
                bestIndex= nextIndex;
                maxDiff= diff;
            }
            ++nextIndex;
        }
    };
}
template<class T> int KDBuilder<T>
::choosePrecise(int nodeIndex,int *beginIDs,int *endIDs,int) const {
    ASSERT( nodeIndex>0 && beginIDs && endIDs && beginIDs<endIDs );
//  temporary storage for computed bounding box
    BoundsPair boundsStorage[length];
    const BoundsPair *localBounds;
//  find out the bounding box
    if ( nodeIndex>1 ) { // compute the bounding box
        localBounds= boundsStorage;
        getBounds( beginIDs, endIDs, boundsStorage );
    } else // we are in the root -> we can use already computed bounds
        localBounds= this->bounds;
//  find and return the longest coordinate
    MaxDiffCoord<T> mdc= for_each
        ( localBounds+1, localBounds+length, MaxDiffCoord<T>(localBounds[0]) );
    ASSERT( mdc.nextIndex == length );
    return mdc.bestIndex;
}

template<class T> int KDBuilder<T>::chooseApprox(int nodeIndex,int*,int*,int) const {
    using namespace FieldMath;
    ASSERT(nodeIndex>0);

    int myDepth= log2ceil(nodeIndex+1)-1;
    if (!myDepth) { // I'm in the root - copy the bounds
        ASSERT(nodeIndex==1);
        chooserTmp= new BoundsPair[length*(depth+1)]; // allocate my temporary bound-array
        assign( bounds, length, chooserTmp );
    }

    Bounds myBounds= chooserTmp+length*myDepth;
    if (myDepth) { // I'm not the root - copy parent's bounds and modify them
        const typename Tree::Node &parent= nodes[nodeIndex/2];
        Bounds parentBounds= myBounds-length;
        if (nodeIndex%2) {  // I'm the right son -> bounds on this level not initialized
            assign( parentBounds, length, myBounds );       // copying parent bounds
            myBounds[parent.coord][0]= parent.threshold;    // adjusting the lower bound
        } else // I'm the left son
            if ( nodeIndex+1 < count ) { // almost the same as brother -> only adjust the coordinate
                myBounds[parent.coord][0]= parentBounds[parent.coord][0];
                myBounds[parent.coord][1]= parent.threshold;
            } else { // I've got no brother
                ASSERT( nodeIndex+1 == count );
                assign( parentBounds, length, myBounds );
                myBounds[parent.coord][1]= parent.threshold;
            }
    }
//  find out the widest dimension
    MaxDiffCoord<T> mdc= for_each( myBounds+1, myBounds+length, MaxDiffCoord<T>(myBounds[0]) );
    ASSERT( mdc.nextIndex == length );
    return mdc.bestIndex;
}

namespace NOSPACE {
    template<class T> class IndexComparator {
        const T *data;  
        int length;     
    public:
        IndexComparator(const T *data_,int length_,int index_)
        : data(data_+index_), length(length_) {}
        bool operator()(int a,int b) const
            { return data[a*length] < data[b*length]; }
    };
}
template<class T> void KDBuilder<T>
::buildNode(int nodeIndex,int *beginIDs,int *endIDs,int depthLeft) {
    int count= endIDs-beginIDs; // owershadowing Tree::count
//  check we've got at least one vector and the depth&count are adequate to each other
    ASSERT( count>=2 && powers[depthLeft-1]<count && count<=powers[depthLeft] );
    --depthLeft;
//  find out where to split - how many items should be on the left to have the heap-shape
    bool shallowRight= ( count <= powers[depthLeft]+powers[depthLeft-1] );
    int *middle= shallowRight
        ? endIDs-powers[depthLeft-1]
        : beginIDs+powers[depthLeft];
//  find out the dividing coordinate and find the "median" in this coordinate
    int coord= (this->*chooser)(nodeIndex,beginIDs,endIDs,depthLeft);
    nth_element( beginIDs, middle , endIDs, IndexComparator<T>(data,length,coord) );
//  fill the node's data (dividing coordinate and its threshold)
    nodes[nodeIndex].coord= coord;
    nodes[nodeIndex].threshold= data[*middle*length+coord]; // min. value of the right son
//  recurse on both halves (if needed; fall-through switch)
    switch (count) {
    default: // we've got enough nodes - build both subtrees (fall through)
    //  build the right subtree
        buildNode( 2*nodeIndex+1, middle, endIDs, depthLeft-shallowRight );
    case 3: // only a pair in the first half
    //  build the left subtree
        buildNode( 2*nodeIndex, beginIDs, middle, depthLeft );
    case 2: // nothing needs to be sorted
        ;
    }
}

#endif // KDTREE_HEADER_

---