Thread-safe refactoring of the class KdTree.

Removed methods:
void setMaxNofNeighbors(unsigned int k);
inline int getNofFoundNeighbors(void);
inline const VectorType& getNeighbor(int i);
inline unsigned int getNeighborId(int i);
inline float getNeighborSquaredDistance(int i);

Added methods:
void doQueryDist(const VectorType& queryPoint, float dist, std::vector<unsigned int>& points, std::vector<Scalar>& sqrareDists);
void doQueryClosest(const VectorType& queryPoint, unsigned int& index, Scalar& dist);

Changed methods:
void doQueryK(const VectorType& queryPoint,  int k, PriorityQueue& mNeighborQueue);

vcg/space/index/kdtree/kdtree.h

@@ -1,14 +1,16 @@
-#ifndef KDTREE_H
+#ifndef KDTREE_VCG_H
-#define KDTREE_H
+#define KDTREE_VCG_H
+
+#include <vcg/space/point3.h>
+#include <vcg/space/box3.h>
+#include <vcg/space/index/kdtree/priorityqueue.h>
 
-#include "../../point3.h"
-#include "../../box3.h"
-#include "mlsutils.h"
-#include "priorityqueue.h"
 #include <vector>
 #include <limits>
 #include <iostream>
 
+namespace vcg {
+
 	template<typename _DataType>
 	class ConstDataWrapper
 	{
@@ -50,7 +52,8 @@ public:
 	};
 
 	/**
- * This class allows to create a Kd-Tree thought to perform the k-nearest neighbour query
+	* This class allows to create a Kd-Tree thought to perform the neighbour query (radius search, knn-nearest serach and closest search). 
+	* The class implemetantion is thread-safe.
 	*/
 	template<typename _Scalar>
 	class KdTree
@@ -61,6 +64,8 @@ public:
 		typedef vcg::Point3<Scalar> VectorType;
 		typedef vcg::Box3<Scalar> AxisAlignedBoxType;
 
+		typedef HeapMaxPriorityQueue<int, Scalar> PriorityQueue;
+
 		struct Node
 		{
 			union {
@@ -84,20 +89,17 @@ public:
 		inline const NodeList& _getNodes(void) { return mNodes; }
 		inline const std::vector<VectorType>& _getPoints(void) { return mPoints; }
 
-
-    void setMaxNofNeighbors(unsigned int k);
-    inline int getNofFoundNeighbors(void) { return mNeighborQueue.getNofElements(); }
-    inline const VectorType& getNeighbor(int i) { return mPoints[ mNeighborQueue.getIndex(i) ]; }
-    inline unsigned int getNeighborId(int i) { return mIndices[mNeighborQueue.getIndex(i)]; }
-    inline float getNeighborSquaredDistance(int i) { return mNeighborQueue.getWeight(i); }
-
 	public:
 
 		KdTree(const ConstDataWrapper<VectorType>& points, unsigned int nofPointsPerCell = 16, unsigned int maxDepth = 64);
 
 		~KdTree();
 
-    void doQueryK(const VectorType& p);
+		void doQueryK(const VectorType& queryPoint,  int k, PriorityQueue& mNeighborQueue);
+
+		void doQueryDist(const VectorType& queryPoint, float dist, std::vector<unsigned int>& points, std::vector<Scalar>& sqrareDists);
+
+		void doQueryClosest(const VectorType& queryPoint, unsigned int& index, Scalar& dist);
 
 	protected:
 
@@ -121,10 +123,7 @@ protected:
 		AxisAlignedBoxType mAABB; //BoundingBox
 		NodeList mNodes; //kd-tree nodes
 		std::vector<VectorType> mPoints; //points read from the input DataWrapper
-        std::vector<int> mIndices; //points indices
+		std::vector<unsigned int> mIndices; //points indices
-
-        HeapMaxPriorityQueue<int,Scalar> mNeighborQueue; //used to perform the knn-query
-        QueryNode mNodeStack[64]; //used in the implementation of the knn-query
 	};
 
 	template<typename Scalar>
@@ -154,11 +153,6 @@ KdTree<Scalar>::~KdTree()
 	{
 	}
 
-template<typename Scalar>
-void KdTree<Scalar>::setMaxNofNeighbors(unsigned int k)
-{
-        mNeighborQueue.setMaxSize(k);
-}
 
 	/** Performs the kNN query.
 	*
@@ -173,15 +167,17 @@ void KdTree<Scalar>::setMaxNofNeighbors(unsigned int k)
 	* But, again, priority queue insertions and deletions are quite involved, and therefore
 	* a simple stack is by far much faster.
 	*
-  * The result of the query, the k-nearest neighbors, are internally stored into a stack, where the
+	* The result of the query, the k-nearest neighbors, are stored into the stack mNeighborQueue, where the
-  * topmost element [0] is NOT the nearest but the farthest!! (they are not sorted but arranged into a heap)
+	* topmost element [0] is NOT the nearest but the farthest!! (they are not sorted but arranged into a heap).
 	*/
 	template<typename Scalar>
-void KdTree<Scalar>::doQueryK(const VectorType& queryPoint)
+	void KdTree<Scalar>::doQueryK(const VectorType& queryPoint, int k, PriorityQueue& mNeighborQueue)
 	{
+		mNeighborQueue.setMaxSize(k);
 		mNeighborQueue.init();
 		mNeighborQueue.insert(0xffffffff, std::numeric_limits<Scalar>::max());
 
+		QueryNode mNodeStack[64];
 		mNodeStack[0].nodeId = 0;
 		mNodeStack[0].sq = 0.f;
 		unsigned int count = 1;
@@ -208,7 +204,7 @@ void KdTree<Scalar>::doQueryK(const VectorType& queryPoint)
 					unsigned int end = node.start+node.size;
 					//adding the element of the leaf to the heap
 					for (unsigned int i=node.start ; i<end ; ++i)
-                                                mNeighborQueue.insert(i, vcg::SquaredNorm(queryPoint - mPoints[i]));
+						mNeighborQueue.insert(mIndices[i], vcg::SquaredNorm(queryPoint - mPoints[i]));
 				}
 				//otherwise, if we're not on a leaf
 				else
@@ -244,6 +240,139 @@ void KdTree<Scalar>::doQueryK(const VectorType& queryPoint)
 		}
 	}
 
+
+	/** Performs the distance query.
+	*
+	* The result of the query, all the points within the distance dist form the query point, is the vector of the indeces
+	* and the vector of the squared distances from the query point.
+	*/
+	template<typename Scalar>
+	void KdTree<Scalar>::doQueryDist(const VectorType& queryPoint, float dist, std::vector<unsigned int>& points, std::vector<Scalar>& sqrareDists)
+	{
+		QueryNode mNodeStack[64];
+		mNodeStack[0].nodeId = 0;
+		mNodeStack[0].sq = 0.f;
+		unsigned int count = 1;
+
+		float sqrareDist = dist*dist;
+		while (count)
+		{
+			QueryNode& qnode = mNodeStack[count-1];
+			Node   & node  = mNodes[qnode.nodeId];
+
+			if (qnode.sq < sqrareDist)
+			{
+				if (node.leaf)
+				{
+					--count; // pop
+					unsigned int end = node.start+node.size;
+					for (unsigned int i=node.start ; i<end ; ++i)
+					{
+						float pointSquareDist = vcg::SquaredNorm(queryPoint - mPoints[i]);
+						if (pointSquareDist < sqrareDist)
+						{
+							points.push_back(mIndices[i]);
+							sqrareDists.push_back(pointSquareDist);
+						}
+					}
+				}
+				else
+				{
+					// replace the stack top by the farthest and push the closest
+					float new_off = queryPoint[node.dim] - node.splitValue;
+					if (new_off < 0.)
+					{
+						mNodeStack[count].nodeId  = node.firstChildId;
+						qnode.nodeId = node.firstChildId+1;
+					}
+					else
+					{
+						mNodeStack[count].nodeId  = node.firstChildId+1;
+						qnode.nodeId = node.firstChildId;
+					}
+					mNodeStack[count].sq = qnode.sq;
+					qnode.sq = new_off*new_off;
+					++count;
+				}
+			}
+			else
+			{
+				// pop
+				--count;
+			}
+		}
+	}
+
+
+	/** Searchs the closest point.
+	*
+	* The result of the query, the closest point to the query point, is the index of the point and 
+	* and the squared distance from the query point.
+	*/
+	template<typename Scalar>
+	void KdTree<Scalar>::doQueryClosest(const VectorType& queryPoint, unsigned int& index, Scalar& dist)
+	{
+		QueryNode mNodeStack[64];
+		mNodeStack[0].nodeId = 0;
+		mNodeStack[0].sq = 0.f;
+		unsigned int count = 1;
+
+		int minIndex = mIndices.size() / 2;
+		Scalar minDist = vcg::SquaredNorm(queryPoint - mPoints[minIndex]);
+		minIndex = mIndices[minIndex];
+
+		while (count)
+		{
+			QueryNode& qnode = mNodeStack[count-1];
+			Node   & node  = mNodes[qnode.nodeId];
+
+			if (qnode.sq < minDist)
+			{
+				if (node.leaf)
+				{
+					--count; // pop
+					unsigned int end = node.start+node.size;
+					for (unsigned int i=node.start ; i<end ; ++i)
+					{
+						float pointSquareDist = vcg::SquaredNorm(queryPoint - mPoints[i]);
+						if (pointSquareDist < minDist)
+						{
+							minDist = pointSquareDist;
+							minIndex = mIndices[i];
+						}
+					}
+				}
+				else
+				{
+					// replace the stack top by the farthest and push the closest
+					float new_off = queryPoint[node.dim] - node.splitValue;
+					if (new_off < 0.)
+					{
+						mNodeStack[count].nodeId  = node.firstChildId;
+						qnode.nodeId = node.firstChildId+1;
+					}
+					else
+					{
+						mNodeStack[count].nodeId  = node.firstChildId+1;
+						qnode.nodeId = node.firstChildId;
+					}
+					mNodeStack[count].sq = qnode.sq;
+					qnode.sq = new_off*new_off;
+					++count;
+				}
+			}
+			else
+			{
+				// pop
+				--count;
+			}
+		}
+		index = minIndex;
+		dist = minDist;
+	}
+
+
+
 	/**
 	* Split the subarray between start and end in two part, one with the elements less than splitValue,
 	* the other with the elements greater or equal than splitValue. The elements are compared
@@ -301,7 +430,12 @@ void KdTree<Scalar>::createTree(unsigned int nodeId, unsigned int start, unsigne
 		VectorType diag = aabb.max - aabb.min;
 
 		//the split "dim" is the dimension of the box with the biggest value
-        unsigned int dim = vcg::MaxCoeffId(diag);
+		unsigned int dim;
+		if (diag.X() > diag.Y())
+			dim = diag.X() > diag.Z() ? 0 : 2;
+		else
+			dim = diag.Y() > diag.Z() ? 1 : 2;
+		
 		node.dim = dim;
 		//we divide the bounding box in 2 partitions, considering the average of the "dim" dimension
 		node.splitValue = Scalar(0.5*(aabb.max[dim] + aabb.min[dim]));
@@ -347,6 +481,6 @@ void KdTree<Scalar>::createTree(unsigned int nodeId, unsigned int start, unsigne
 			}
 		}
 	}
+}
 
 #endif
-