KDTreeLinkerAlgo.h

#ifndef KDTreeLinkerAlgoTemplated_h
#define KDTreeLinkerAlgoTemplated_h

#include "KDTreeLinkerTools.h"

#include <cassert>
#include <vector>

// Class that implements the KDTree partition of 2D space and 
// a closest point search algorithme.

template <typename DATA>
class KDTreeLinkerAlgo
{
 public:
  KDTreeLinkerAlgo();
  
  // Dtor calls clear()
  ~KDTreeLinkerAlgo();
  
  // Here we build the KD tree from the "eltList" in the space define by "region".
  void build(std::vector<KDTreeNodeInfo<DATA> > 	&eltList,
	     const KDTreeBox				&region);
  
  // Here we search in the KDTree for all points that would be 
  // contained in the given searchbox. The founded points are stored in resRecHitList.
  void search(const KDTreeBox				&searchBox,
	      std::vector<DATA>	&resRecHitList);

  // This reurns true if the tree is empty
  bool empty() {return nodePool_.empty();}

  // This returns the number of nodes + leaves in the tree
  // (nElements should be (size() +1)/2)
  int size() { return nodePool_.size();}

  // This method clears all allocated structures.
  void clear();
  
 private:
  // The node pool allow us to do just 1 call to new for each tree building.
  KDTreeNodes<DATA> nodePool_;
  
  std::vector<DATA>	*closestNeighbour;
  std::vector<KDTreeNodeInfo<DATA> >	*initialEltList;
  
 private:
 
  //Fast median search with Wirth algorithm in eltList between low and high indexes.
  int medianSearch(int					low,
		   int					high,
		   int					treeDepth);

  // Recursif kdtree builder. Is called by build()
  int recBuild(int				low,
               int				hight,
               int				depth);

  // Recursif kdtree search. Is called by search()
  void recSearch(int			current,
                 VALUE dimCurrMin, VALUE dimCurrMax,
                 VALUE dimOtherMin, VALUE dimOtherMax);

  // This method frees the KDTree.     
  void clearTree();
};


//Implementation

template < typename DATA >
void
KDTreeLinkerAlgo<DATA>::build(std::vector<KDTreeNodeInfo<DATA> >	&eltList, 
			      const KDTreeBox				&region)
{
  if (eltList.size()) {
    initialEltList = &eltList;
    
    size_t size = initialEltList->size();
    nodePool_.build(size);
    
    // Here we build the KDTree
    int root = recBuild(0, size, 0);
    assert(root == 0);
    
    initialEltList = 0;
  }
}
 
//Fast median search with Wirth algorithm in eltList between low and high indexes.
template < typename DATA >
int
KDTreeLinkerAlgo<DATA>::medianSearch(int	low,
				     int	high,
				     int	treeDepth)
{
  int nbrElts = high - low;
  int median = (nbrElts & 1)	? nbrElts / 2 
				: nbrElts / 2 - 1;
  median += low;

  int l = low;
  int m = high - 1;
  
  while (l < m) {
    KDTreeNodeInfo<DATA> elt = (*initialEltList)[median];
    int i = l;
    int j = m;

    do {
      // The even depth is associated to dim1 dimension
      // The odd one to dim2 dimension
      if (treeDepth & 1) {
	while ((*initialEltList)[i].dim[1] < elt.dim[1]) i++;
	while ((*initialEltList)[j].dim[1] > elt.dim[1]) j--;
      } else {
	while ((*initialEltList)[i].dim[0] < elt.dim[0]) i++;
	while ((*initialEltList)[j].dim[0] > elt.dim[0]) j--;
      }

      if (i <= j){
	std::swap((*initialEltList)[i], (*initialEltList)[j]);
	i++; 
	j--;
      }
    } while (i <= j);
    if (j < median) l = i;
    if (i > median) m = j;
  }

  return median;
}



template < typename DATA >
void
KDTreeLinkerAlgo<DATA>::search(const KDTreeBox		&trackBox,
			 std::vector<DATA> &recHits)
{
  if (!empty()) {
    closestNeighbour = &recHits;
    recSearch(0, trackBox.dim1min, trackBox.dim1max, trackBox.dim2min, trackBox.dim2max);
    closestNeighbour = 0;
  }
}


template < typename DATA >
void 
KDTreeLinkerAlgo<DATA>::recSearch(int	current,
                                  VALUE dimCurrMin, VALUE dimCurrMax,
                                  VALUE dimOtherMin, VALUE dimOtherMax)
{
  // Iterate until leaf is found, or there are no children in the
  // search window. If search has to proceed on both children, proceed
  // the search to left child via recursion. Swap search window
  // dimension on alternate levels.
  while(true) {
    int right = nodePool_.right[current];
    if(nodePool_.isLeaf(right)) {
      VALUE dimCurr = nodePool_.median[current];

      // If point inside the rectangle/area
      // Use intentionally bit-wise & instead of logical && for better
      // performance. It is faster to always do all comparisons than to
      // allow use of branches to not do some if any of the first ones
      // is false.
      if((dimCurr >= dimCurrMin) & (dimCurr <= dimCurrMax)) {
        VALUE dimOther = nodePool_.dimOther[current];
        if((dimOther >= dimOtherMin) & (dimOther <= dimOtherMax)) {
          closestNeighbour->push_back(nodePool_.data[current]);
        }
      }
      break;
    }
    else {
      VALUE median = nodePool_.median[current];

      bool goLeft = (dimCurrMin <= median);
      bool goRight = (dimCurrMax >= median);

      // Swap dimension for the next search level
      std::swap(dimCurrMin, dimOtherMin);
      std::swap(dimCurrMax, dimOtherMax);
      if(goLeft & goRight) {
        int left = current+1;
        recSearch(left, dimCurrMin, dimCurrMax, dimOtherMin, dimOtherMax);
        // continue with right
        current = right;
      }
      else if(goLeft) {
        ++current;
      }
      else if(goRight) {
        current = right;
      }
      else {
        break;
      }
    }
  }
}

template <typename DATA>
KDTreeLinkerAlgo<DATA>::KDTreeLinkerAlgo()
{
}

template <typename DATA>
KDTreeLinkerAlgo<DATA>::~KDTreeLinkerAlgo()
{
  clear();
}


template <typename DATA>
void 
KDTreeLinkerAlgo<DATA>::clearTree()
{
  nodePool_.clear();
}

template <typename DATA>
void 
KDTreeLinkerAlgo<DATA>::clear()
{
  clearTree();
}


template <typename DATA>
int
KDTreeLinkerAlgo<DATA>::recBuild(int					low, 
                                 int					high, 
                                 int					depth)
{
  int portionSize = high - low;
  int dimIndex = depth&1;

  if (portionSize == 1) { // Leaf case
    int leaf = nodePool_.getNextNode();
    const KDTreeNodeInfo<DATA>& info = (*initialEltList)[low];
    nodePool_.right[leaf] = 0;
    nodePool_.median[leaf] = info.dim[dimIndex]; // dimCurrent
    nodePool_.dimOther[leaf] = info.dim[1-dimIndex];
    nodePool_.data[leaf] = info.data;
    return leaf;

  } else { // Node case
    
    // The even depth is associated to dim1 dimension
    // The odd one to dim2 dimension
    int medianId = medianSearch(low, high, depth);
    VALUE medianVal = (*initialEltList)[medianId].dim[dimIndex];

    // We create the node
    int nodeInd = nodePool_.getNextNode();
    nodePool_.median[nodeInd] = medianVal;

    ++depth;
    ++medianId;

    // We recursively build the son nodes
    int left = recBuild(low, medianId, depth);
    assert(nodeInd+1 == left);
    nodePool_.right[nodeInd] = recBuild(medianId, high, depth);

    return nodeInd;
  }
}

#endif