BallTree.h

#ifndef GEOMETRY_BallTREE_H
#define GEOMETRY_BallTREE_H

#include "KDTree.h"
#include <KrisLibrary/math/vector.h>
#include <vector>
#include <functional>
#include <memory>
using namespace Math;

namespace Geometry {

class BallTreeNode
{
public:
  typedef KDTree::Point Point;

  BallTreeNode();

  inline bool IsLeaf() const { return children.empty(); }
  int MaxDepth() const;
  int MinDepth() const;

  int MaxLeafSize() const;
  int MinLeafSize() const;

  int TreeSize() const;

  Vector center;
  Real radius;

  std::vector<Point> pts;

  BallTreeNode* parent;
  std::vector<std::unique_ptr<BallTreeNode> > children;
};

class BallTree
{
public:
  typedef std::function<double(const Vector&,const Vector&)> Metric;
  typedef std::function<void(const Vector&,const Vector&,Real,Vector&)> Interpolator;

  Metric metric;
  BallTreeNode root;
  Interpolator interpolator;
  bool cartesian;
  int splitsPerNode;

  BallTree(Metric metric,int numSplitsPerNode=2);
  BallTree(Metric metric,Interpolator interpolator,int numSplitsPerNode=2);

  inline int MaxDepth() const { return root.MaxDepth(); }
  inline int MinDepth() const { return root.MinDepth(); }
  inline int MaxLeafSize() const { return root.MaxLeafSize(); }
  inline int MinLeafSize() const { return root.MinLeafSize(); }

  inline int TreeSize() const { return root.TreeSize(); }

  void Build(const std::vector<Vector>& p,int maxDepth);

  BallTreeNode* Insert(const Vector& p,int id,int maxLeafPoints=-1);

  bool Split(BallTreeNode* node);
  void Join(BallTreeNode* node);
  void Fit(BallTreeNode* node,bool tight);
  void Clear();

  int ClosestPoint(const Vector& pt,Real& dist) const;

  int PointWithin(const Vector& pt,Real& dist) const;

  void ClosePoints(const Vector& pt,Real radius,std::vector<Real>& distances,std::vector<int>& ids) const;

  void KClosestPoints(const Vector& pt,int k,Real* dist,int* idx) const;

private:
  void _ClosestPoint(const BallTreeNode* node,const Vector& pt,Real& dist,int& idx) const;
  void _ClosePoints(const BallTreeNode* node,const Vector& pt,Real radius,std::vector<Real>& distances,std::vector<int>& ids) const;
  void _KClosestPoints(const BallTreeNode* node,const Vector& pt,int k,Real* dist,int* idx,int& maxdist) const;
  BallTreeNode* _LookupClosestLeaf(BallTreeNode* node,const Vector& pt,Real& dist);
};


} //namespace Geometry

#endif