untested kd tree code

kd tree

#ifndef COMMON_KDTREE_H
#define COMMON_KDTREE_H

/*
 * Copyright (c) 2015 Daniel Perry
 *
 */

#include <vector>
using std::vector;

#include <boost/shared_ptr.hpp>
using boost::shared_ptr;

namespace common
{
namespace kdtree
{

/** Simple D-dim spatial axis-aligned region class.
 */
template<int DIM>
class Region
{
	const int D = DIM;
	float lower[DIM];
	float upper[DIM];

	/** default ctor
	 */
	Region()
	{
		for(int i=0; i<D; ++i)
		{
			lower[i] = 0;
			upper[i] = 0;
		}
	}
	/** ctor
	 */
	Region(float * low, float * up)
	{
		for(int i=0; i<D; ++i)
		{
			lower[i] = low[i];
			upper[i] = up[i];
		}
	}
	/** copy ctor
	 */
	Region(const Region & reg)
	{
		for(int i=0; i<D; ++i)
		{
			lower[i] = reg.lower[i];
			upper[i] = reg.upper[i];
		}
	}
	/** assignment operator
	 */
	Region & operator=(const Region & reg)
	{
		for(int i=0; i<D; ++i)
		{
			lower[i] = reg.lower[i];
			upper[i] = reg.upper[i];
		}
		return (*this);
	}

	/** Size along a specific dimension.
	 */
	int Size(int dim)
	{
		return upper[dim] - lower[dim];
	}

	/** test if a point is contained in the region.
	 */
	template<class Point>
	bool Contains(const Point & pt)
	{
		for(int i=0; i<D; ++i)
		{
			if(pt[i] < lower[i] || upper[i] < pt[i]) return false;
		}
		return true;
	}

	/** Halve the region in one dimension
	 * @param dim - dimension along which to halve.
	 * @param lower - return the region less than the middle value.
	 *
	 * RETURN:
	 * @param value - return the value along which the division is made.
	 * @return the new region.
	 */
	Region Halve(int dim, bool lower, float & value)
	{
		Region region(*this);
		value = (upper[dim] - lower[dim]) / 2.f;
		if(lower)
		{
			region.upper[dim] = value;
		}
		else
		{
			region.lower[dim] = value;
		}
		return region;
	}
};

/** Unidirectional Node for a binary tree in d-dimensional space
 */
 template <int DIM, class Cargo>
class Node
{
public:
	const int D = DIM;
	typedef Node<Cargo,D> NodeT;
	/** pointers to child nodes. */
	shared_ptr<NodeT> child[2];
	/** spatial region enclosed by this node. */
	Region region;
	/** which dimension to divide along for children. */
	int dim;
	/** value on which to divide for children. */
	float value;

	/** leaf nodes point to a list of cargo contained within the space. */
	vector<Cargo*> cargo;

	/** ctor 
	 */
	Node(int dim, const Region & reg)
	:child(), // no children.
	 parent(0), // null
	 region(reg), 
	 dim(dim), 
	 value(0) 
	{
		// init with no chil'en, set points to null:
		child[0].reset(0);
		child[1].reset(0);
	}

	/** Spawn children until the smallest dimensions of the kd-tree would be smaller than space_limit if divided.
		* @param space_limit - precision of the leaf division.
	 */
	void Spawn(float space_limit)
	{
		int new_dim = (dim + 1)%D;
		for(int i=0; i<2; ++i)
		{
			Region child_region = region.Halve(new_dim, !i, value);

			bool small_enough = true;
			for(int i=0; i<D; ++i)
			{
				if(child_region.Size(i) < space_limit)
				{
					small_enough = false;
					break;
				}
			}
			if(small_enough) return;

			child[i].reset( new Node(new_dim, child_region) );
			child[i]->Spawn(space_limit); // recursive call to children.
		}
	}

	/** Locate a point within the kd-tree
		* @param pt - the point to find.
		* @return the leaf node in which it resides.
	 */
	template<class Point>
	Node & Locate(const Point & pt)
	{
		if(!child[0]) // leaf node!
		{
			return (*this);
		}

		// keep looking recursively:
		if(pt[dim] < value)
		{
			child[0]->Locate(pt);
		}
		else
		{
			child[1]->Locate(pt);
		}
	}
	
	/** Insert some cargo associated with a point into the kd-tree
	  * @param pt - spatial location associated with cargo.
		* @param c - the cargo to insert.
	 */
	template<class Point>
	void Insert(const Point & pt, Cargo * c)
	{
		Node & leaf = Locate(pt);
		leaf.cargo.push_back(c);
	}

	/** Insert some cargo associated with a region into the kd-tree
	  * NOTE: this cargo may end up in multiple nodes.
		*
	  * @param reg - spatial region associated with cargo.
		* @param c - the cargo to insert.
	 */
	void Insert(const Region & reg, Cargo * c)
	{
		if(!child[0]) // leaf node!
		{
			cargo.push_back(c);
			return;
		}

		// keep looking recursively:
		if(reg.lower[dim] < value)
		{
			child[0]->Insert(reg,c);
		}
		if(reg.upper[dim] >= value)
		{
			child[1]->Insert(reg,c);
		}
	}
};


} // end namespace
} // end namespace

#endif