Octree.h

#ifndef GEOMETRY_OCTREE_H
#define GEOMETRY_OCTREE_H

#include <KrisLibrary/math3d/AABB3D.h>
#include <KrisLibrary/math3d/Box3D.h>
#include <KrisLibrary/math3d/Ray3D.h>
#include <vector>
#include <list>

namespace Geometry {
  using namespace std;
  using namespace Math3D;


struct OctreeNode
{
  //child indices
  enum {LLL,LLU,LUL,LUU,ULL,ULU,UUL,UUU};

  AABB3D bb;
  int parentIndex;
  int childIndices[8];
};

class Octree
{
 public:
  Octree(const AABB3D& bb);
  virtual ~Octree() {}
  int NumNodes() const { return (int)nodes.size()-(int)freeNodes.size(); }
  int Size() const { return NumNodes(); }
  bool IsLeaf(const OctreeNode& n) const { return n.childIndices[0] < 0; }
  int Depth(const OctreeNode& n) const { if(n.parentIndex<0) return 0; return 1+Depth(nodes[n.parentIndex]); }
  int Index(const OctreeNode& n) const { return int(&n - &nodes[0]); }
  const OctreeNode& Node(int index) const { return nodes[index]; }
  OctreeNode& Node(int index) { return nodes[index]; }
  int MaxDepth() const;
  void SplitToDepth(int d);
  void SplitToDepth(OctreeNode& n,int d);
  OctreeNode* SplitToDepth(OctreeNode& n,const Vector3& point,int d);
  void SplitToResolution(const Vector3& res);
  void SplitToResolution(OctreeNode& n,const Vector3& res);
  OctreeNode* SplitToResolution(OctreeNode& n,const Vector3& point,const Vector3& res);
  int Child(const OctreeNode& n,const Vector3& pt) const;
  void Range(const OctreeNode& n,int child,AABB3D& bbchild) const;
  OctreeNode* Lookup(const Vector3& point);
  OctreeNode* Lookup(OctreeNode& root,const Vector3& point);
  OctreeNode* Lookup(OctreeNode& root,const Vector3& point,int depthMax);
  void BoxLookup(const Vector3& bmin,const Vector3& bmax,vector<int>& nodeIndices) const;
  void BoxLookup(const Box3D& b,vector<int>& nodeIndices) const;
  void BallLookup(const Vector3& c,Real r,vector<int>& nodeIndices) const;
  void RayLookup(const Ray3D& ray,vector<int>& nodeindices) const;
  void FattenedRayLookup(const Ray3D& ray,Real radius,vector<int>& nodeindices) const;
  virtual void Split(int nodeindex);
  virtual void Join(int nodeindex);

 protected:
  virtual int AddNode(int parent=-1);
  virtual void DeleteNode(int);
  void _RayLookup(int nodeindex,const Ray3D& ray,vector<int>& nodeindices) const;
  void _FattenedRayLookup(int nodeindex,const Ray3D& ray,Real radius,vector<int>& nodeindices) const;

  vector<OctreeNode> nodes;
  list<int> freeNodes;
};

class OctreePointSet : public Octree
{
 public:
  OctreePointSet(const AABB3D& bbox,int maxPointsPerCell=1,Real minCellSize=0);
  virtual ~OctreePointSet() {}
  void GetPoints(int node,vector<Vector3>& pts) const;
  void GetPoints(const OctreeNode& node,vector<Vector3>& pts) const { GetPoints(Index(node),pts); }
  void GetPointIDs(int node,vector<int>& ids) const;
  void GetPointIDs(const OctreeNode& node,vector<int>& ids) const { GetPointIDs(Index(node),ids); }
  void Add(const Vector3& pt,int id=-1);
  void BoxQuery(const Vector3& bmin,const Vector3& bmax,vector<Vector3>& points,vector<int>& ids) const;
  void BoxQuery(const Box3D& b,vector<Vector3>& points,vector<int>& ids) const;
  void BallQuery(const Vector3& c,Real r,vector<Vector3>& points,vector<int>& ids) const;
  void RayQuery(const Ray3D& r,Real radius,vector<Vector3>& points,vector<int>& ids) const;
  int RayCast(const Ray3D& r,Real radius) const;
  bool NearestNeighbor(const Vector3& c,Vector3& closest,int& id) const;
  void KNearestNeighbors(const Vector3& c,int k,vector<Vector3>& closest,vector<int>& ids) const;
  void Collapse(int maxSize=0);
  void FitToPoints();
  const Sphere3D& Ball(int index) const { return balls[index]; }

 protected:
  Real _NearestNeighbor(const OctreeNode& n,const Vector3& c,Vector3& closest,int& id,Real minDist) const;
  int _KNearestNeighbors(const OctreeNode& n,const Vector3& c,vector<Vector3>& closest,vector<int>& ids,vector<Real>& distances,int kmin) const;
  virtual int AddNode(int parent=-1);
  virtual void DeleteNode(int id);
  virtual void Split(int nodeindex);
  virtual void Join(int nodeindex);

  int maxPointsPerCell;
  Real minCellSize;
  vector<vector<int> > indexLists;
  vector<Vector3> points;
  vector<int> ids;
  vector<Sphere3D> balls;
  bool fit;
};

class OctreeScalarField : public Octree
{
 public:
  OctreeScalarField(const AABB3D& bb,Real defaultValue = -Inf);
  virtual ~OctreeScalarField() {}
  void Set(const Vector3& pt,Real value,int id=-1);
  Real Value(const Vector3& pt) const;
  bool ValueIn(const Vector3& pt,Real valueMin,Real valueMax) const;
  bool ValueGreater(const Vector3& pt,Real bound) const;
  bool ValueLess(const Vector3& pt,Real bound) const;
  void BoxLookupRange(const Vector3& bmin,const Vector3& bmax,Real valueMin,Real valueMax,vector<int>& nodeIndices,bool inclusive=true) const;
  void BoxLookupGreater(const Vector3& bmin,const Vector3& bmax,Real bound,vector<int>& nodeIndices) const;
  void BoxLookupLess(const Vector3& bmin,const Vector3& bmax,Real bound,vector<int>& nodeIndices) const;
  void Collapse(Real tolerance=0);

 protected:
  struct Data {
    Real value;
    Real valueMin,valueMax;
    int id;
  };

  virtual int AddNode(int parent=-1);
  virtual void DeleteNode(int id);
  virtual void Split(int nodeindex);
  virtual void Join(int nodeindex);

  Real defaultValue;
  vector<Data> data;
};

} //namespcae Geometry

#endif