kdTree.h

Go to the documentation of this file.

#ifndef DLR_COMPUTERVISION_KDTREE_H
#define DLR_COMPUTERVISION_KDTREE_H

#include <functional>
#include <vector>


namespace dlr {

  namespace computerVision {

    template <unsigned int Dimension, class Type>
    class KDComparator
      : public std::binary_function<Type, Type, bool>
    {
      public:
      
      KDComparator(unsigned int axis = 0)
        : m_axis(axis) {};
      
      
      virtual
      ~KDComparator() {}
      

      double
      computeDistance(Type const& arg0, Type const& arg1) const {
        double distance = 0.0;
        for(size_t ii = 0; ii < Dimension; ++ii) {
          double newTerm = arg0[ii] - arg1[ii];
          distance += newTerm * newTerm;
        }
        return distance;
      }


      double
      getPrimarySeparation(Type const& arg0, Type const& arg1) const {
        double difference = arg0[m_axis] - arg1[m_axis];
        return difference * difference;
      }


      bool isEqual(Type const& arg0, Type const& arg1) const {
        return arg0 == arg1;
      }
      
        
      bool
      operator()(Type const& arg0, Type const& arg1) const {
        if(arg0[m_axis] < arg1[m_axis]) {
          return true;
        }
        if(arg0[m_axis] > arg1[m_axis]) {
          return false;
        }
        for(size_t ii = m_axis + 1; ii < Dimension; ++ii) {
          if(arg0[ii] < arg1[ii]) {
            return true;
          }
          if(arg0[ii] > arg1[ii]) {
            return false;
          }
        }
        for(size_t ii = 0; ii < m_axis; ++ii) {
          if(arg0[ii] < arg1[ii]) {
            return true;
          }
          if(arg0[ii] > arg1[ii]) {
            return false;
          }
        }
        return false;
      }
      
      
    private:

      unsigned int m_axis;
      
    };


    
    template <unsigned int Dimension, class Type>
    class KDTree {
    public:

      KDTree();


      template <class Iter>
      KDTree(Iter beginIter, Iter endIter);


      virtual
      ~KDTree();


      bool
      find(Type const& point) const;


      Type const&
      findNearest(Type const& point, double& distance) const;


      // void
      // rebalance();
      
      
    protected:

      template <class Iter>
      KDTree(Iter beginIter, Iter endIter, size_t vectorSize, size_t level);


      template <class Iter>
      void
      construct(Iter beginIter, Iter endIter, size_t vectorSize, size_t level);

      void
      findNearestIterative(Type const& point,
                           Type const*& bestPointPtr,
                           double& bestDistance) const;

      
      void
      findNearestRecursive(Type const& point,
                           Type const*& bestPointPtr,
                           double& bestDistance) const;
      

      KDComparator<Dimension, Type> m_comparator;
      Type m_point;
      KDTree* m_leftChild;
      KDTree* m_rightChild;
    };

  } // namespace computerVision
  
} // namespace dlr


/* ============ Definitions of inline & template functions ============ */


#include <algorithm>
#include <cmath>
#include <limits>
#include <stack>

namespace dlr {

  namespace computerVision {


    template <unsigned int Dimension, class Type>
    KDTree<Dimension, Type>::
    KDTree()
      : m_comparator(0),
        m_point(),
        m_leftChild(0),
        m_rightChild(0)
    {}


    template <unsigned int Dimension, class Type>
    template <class Iter>
    KDTree<Dimension, Type>::
    KDTree(Iter beginIter, Iter endIter)
      : m_comparator(0),
        m_point(),
        m_leftChild(0),
        m_rightChild(0)
    {
      if(beginIter == endIter) {
        return;
      }

      std::vector<Type> pointVector;
      std::copy(beginIter, endIter, std::back_inserter(pointVector));

      this->construct(pointVector.begin(), pointVector.end(),
                      pointVector.size(), 0);
    }
      
    
    // The destructor cleans up any system resources during destruction.
    template <unsigned int Dimension, class Type>
    KDTree<Dimension, Type>::
    ~KDTree() {
      if(m_leftChild != 0) {
        delete m_leftChild;
      }
      if(m_rightChild != 0) {
        delete m_rightChild;
      }
    }


    template <unsigned int Dimension, class Type>
    bool
    KDTree<Dimension, Type>::
    find(Type const& point) const
    {
      if(m_comparator.isEqual(m_point, point)) {
        return true;
      }

      if(m_comparator(point, m_point)) {
        if(m_leftChild == 0) {
          return false;
        }
        return m_leftChild->find(point);
      } else {
        if(m_rightChild == 0) {
          return false;
        }
        return m_rightChild->find(point);
      }
    }
    

    template <unsigned int Dimension, class Type>
    Type const&
    KDTree<Dimension, Type>::
    findNearest(Type const& point, double& distance) const
    {
      distance = std::numeric_limits<double>::max();
      Type const* bestPointPtr = &m_point;
      // this->findNearestRecursive(point, bestPointPtr, distance);
      this->findNearestIterative(point, bestPointPtr, distance);
      return *bestPointPtr;
    }


    /* ================ Protected ================= */

    template <unsigned int Dimension, class Type>
    template <class Iter>
    KDTree<Dimension, Type>::
    KDTree(Iter beginIter, Iter endIter, size_t vectorSize, size_t level)
      : m_comparator(level % Dimension),
        m_point(),
        m_leftChild(0),
        m_rightChild(0)
    {
      this->construct(beginIter, endIter, vectorSize, level);
    }


    template <unsigned int Dimension, class Type>
    template <class Iter>
    void
    KDTree<Dimension, Type>::
    construct(Iter beginIter, Iter endIter, size_t vectorSize, size_t level)
    {
      std::sort(beginIter, endIter, m_comparator);
      size_t partitionIndex = vectorSize / 2;

      m_point = *(beginIter + partitionIndex);

      if(vectorSize == 1) {
        return;
      }
      m_leftChild = new KDTree(
        beginIter, beginIter + partitionIndex,
        partitionIndex, level + 1);

      if(vectorSize == 2) {
        return;
      }
      m_rightChild = new KDTree(
        beginIter + (partitionIndex + 1), endIter,
        vectorSize - (partitionIndex + 1), level + 1);
    }


    template <unsigned int Dimension, class Type>
    void
    KDTree<Dimension, Type>::
    findNearestIterative(Type const& point,
                         Type const*& bestPointPtr,
                         double& bestDistance) const
    {
      // Contents of this stack are std::pairs in which the first
      // element points to an un-searched tree, and the second element
      // is a lower bound on the distance from argument point to the
      // points contained in the un-searched tree.
      std::stack< std::pair< KDTree<Dimension, Type> const*, double> >
        kdTreeStack;
      kdTreeStack.push(std::make_pair(this, 0.0));

      while(!(kdTreeStack.empty())) {
        KDTree<Dimension, Type> const* currentTree = kdTreeStack.top().first;
        double bound = kdTreeStack.top().second;
        kdTreeStack.pop();

        if(bestDistance < bound) {
          continue;
        }
        
        double myDistance = currentTree->m_comparator.computeDistance(
          point, currentTree->m_point);
        if(myDistance < bestDistance) {
          bestDistance = myDistance;
          bestPointPtr = &currentTree->m_point;
        }

        if(currentTree->m_leftChild == 0 && currentTree->m_rightChild == 0) {
          continue;
        }
      
        KDTree* nearChildPtr;
        KDTree* farChildPtr;
        bool isLeft = currentTree->m_comparator(point, currentTree->m_point);
        if(isLeft) {
          nearChildPtr = currentTree->m_leftChild;
          farChildPtr = currentTree->m_rightChild;
        } else {
          nearChildPtr = currentTree->m_rightChild;
          farChildPtr = currentTree->m_leftChild;
        }

        // Push remote child first, and near child second, so that
        // near child will be popped first, increasing the chance that
        // remote child will be eliminated.

        // Only push the remote child if it's plausible that it
        // contains a closer point than our best so far.  This
        // duplicates a similar test above. Not sure if it's worth
        // duplicating the test to avoid the occasional extra
        // push/pop.
        if(farChildPtr) {
          double remoteDistanceLowerBound =
            currentTree->m_comparator.getPrimarySeparation(
              point, currentTree->m_point);
          if(remoteDistanceLowerBound < bestDistance) {
            kdTreeStack.push(
              std::make_pair(farChildPtr, remoteDistanceLowerBound));
          }
        }

        if(nearChildPtr) {
          kdTreeStack.push(std::make_pair(nearChildPtr, 0.0));
        }
      }
    }

      
    template <unsigned int Dimension, class Type>
    void
    KDTree<Dimension, Type>::
    findNearestRecursive(Type const& point,
                         Type const*& bestPointPtr,
                         double& bestDistance) const
    {
      double myDistance = m_comparator.computeDistance(point, m_point);
      if(myDistance < bestDistance) {
        bestDistance = myDistance;
        bestPointPtr = &m_point;
      }

      if(m_leftChild == 0 && m_rightChild == 0) {
        return;
      }
      
      bool isLeft = m_comparator(point, m_point);
      double remoteDistanceLowerBound = m_comparator.getPrimarySeparation(
        point, m_point);

      KDTree* nearChildPtr;
      KDTree* farChildPtr;
      if(isLeft) {
        nearChildPtr = m_leftChild;
        farChildPtr = m_rightChild;
      } else {
        nearChildPtr = m_rightChild;
        farChildPtr = m_leftChild;
      }        

      if(nearChildPtr) {
        nearChildPtr->findNearestRecursive(point, bestPointPtr, bestDistance);
      }

      if((remoteDistanceLowerBound < bestDistance) && farChildPtr) {
        farChildPtr->findNearestRecursive(point, bestPointPtr, bestDistance);
      }
    }
    
  } // namespace computerVision
  
} // namespace dlr

#endif /* #ifndef DLR_COMPUTERVISION_KDTREE_H */