KDTree.h

#ifndef GEOMETRY_KDTREE_H
#define GEOMETRY_KDTREE_H

#include <KrisLibrary/math/vector.h>
#include <vector>
using namespace Math;

namespace Geometry {

class KDTree
{
 public:
  struct Point {
    Point();
    Point(const Point& pt);
    const Point& operator = (const Point&);
    Vector pt;
    int id;
  };

  static KDTree* Create(const std::vector<Vector>& p, int k, int maxDepth);

  KDTree();
  KDTree(const std::vector<Point>& pts, int k, int depth=0, int maxDepth=100);
  ~KDTree();

  KDTree(const KDTree&) = delete;
  const KDTree& operator=(const KDTree&) const = delete;

  inline bool IsLeaf() const { return splitDim==-1; }
  int MaxDepth() const;
  int MinDepth() const;

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

  int TreeSize() const;

  KDTree* Insert(const Vector& p,int id,int maxLeafPoints=2);
  KDTree* Locate(const Vector& p);

  bool Split(int dimension=-1);
  void Join();
  void Clear();

  int ClosestPoint(const Vector& pt,Real& dist) const;
  int ClosestPoint(const Vector& pt,Real n,const Vector& w,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 ClosePoints(const Vector& pt,Real radius,Real n,const Vector& w,std::vector<Real>& distances,std::vector<int>& ids) const;

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

  static Real Select(const std::vector<Point>& S, int d, int k);

private:
  void _ClosestPoint(const Vector& pt,Real& dist,int& idx) const;
  void _KClosestPoints(const Vector& pt,int k,Real* dist,int* idx,int& maxdist) const;
  void _ClosestPoint2(const Vector& pt,Real& dist,int& idx,Real norm,const Vector& weights) const;
  void _KClosestPoints2(const Vector& pt,int k,Real* dist,int* idx,int& maxdist,Real norm,const Vector& weights) const;

  int depth;
  int splitDim;
  Real splitVal;
  KDTree *pos,*neg;
  std::vector<Point> pts;

  //temporary -- used for performance testing
  int visits;
};

} //namespace Geometry

#endif