KDTree.h

Go to the documentation of this file.

/* Copyright (C) 2004
   K-Dimension Tree Data Structure Implementation

   This library is free software; you can redistribute it and/or
   modify it under the terms of the GNU Lesser General Public
   License as published by the Free Software Foundation; either
   version 2.1 of the License, or (at your option) any later version.

   This library is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
   Lesser General Public License for more details.

   You should have received a copy of the GNU Lesser General Public
   License along with this library; if not, write to the Free Software
   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
*/

#ifndef _KDTREE_H_
#define _KDTREE_H_

#include "Object.h"
#include <stdio.h>
#include <math.h>
#include "KDNode.h"
#include "KDPoint.h"
#include "KDPQueue.h"

namespace RobotFlow {
// Template class for node values
// IMPORTANT: class requires =, -, >, < operators
//            it should usually be either double, float, int, ...
template <class KDData>
class KDTree : public FD::Object
{
public:
        /*
                KDTree

                Default constructor
        */
        KDTree()
        : m_bboxLow(), m_bboxHigh()
        {

        }

        /*
                KDTree

                Other constructor wih user data

                PARAMETERS
                i_dataPts : array of data describing the points used to built the tree.
                                        The array should be constructed as follow:
                                        i_dataPts[<Number of points>][<Dimension size of the points>]
                i_numPts : number of points in the given array
                i_dimSize : dimension size of the points in the given array
                i_bucketSize : bucket size of the leaf nodes i.e. the maximum number of
                                                points strored in a leaf node.
        */
        KDTree(KDData **i_dataPts, int i_numPts, int i_dimSize, int i_bucketSize)
        : m_numPts(i_numPts), m_dimSize(i_dimSize), m_bucketSize(i_bucketSize),
        m_bboxLow(m_dimSize, (KDData)0), m_bboxHigh(m_dimSize, (KDData)0)
        {
                m_root = NULL;
                int *ptsIdx = new int[m_numPts];

                if(m_numPts == 0) {
                        return;
                }

                // First assign identity indexes
                for(int i=0; i<m_numPts; ++i) {
                        ptsIdx[i] = i;
                }

                SetBBoxes(i_dataPts, ptsIdx);

                m_root = RecurBuildTree(i_dataPts, ptsIdx, m_numPts);

                delete [] ptsIdx;
        }

        /*
                ~KDTree

                Destructor
        */
        ~KDTree()
        {
                delete m_root;
        }

        /*
                printOn

                Object routine required to print its content

                PARAMETERS
                out : output stream to print the object content
        */
        void printOn(std::ostream &out) const
        {
                out << "<KDTree " << std::endl;
                out << "<NumPts " << m_numPts << " >" << std::endl;
                out << "<DimSize " << m_dimSize << " >" << std::endl;
                out << "<BucketSize " << m_bucketSize << " >" << std::endl;

                out << "<BBoxLow " << std::endl;
                m_bboxLow.printOn(out);
                out << " >" << std::endl;

                out << "<BBoxHigh " << std::endl;
                m_bboxHigh.printOn(out);
                out << " >" << std::endl;

                if (m_root != NULL) {
                        out << "<KDRoot " << std::endl;
                        m_root->printOn(out);
                        out << " >" << std::endl;
                }

                out << " >" << std::endl;
        }

        /*
                readFrom

                Object routine required to read an Object content

                PARAMETERS
                in : input stream to read the object content
        */
        void readFrom(std::istream &in)
        {
                std::string tag;

                while (1) {
                        char ch;
                        in >> ch;

                        if (ch == '>') {
                                break;
                        }
                        else if (ch != '<') {
                                throw new FD::GeneralException ("KDTree::readFrom : Parse error: '<' expected",__FILE__,__LINE__);
                        }

                        in >> tag;

                        if (tag == "KDTree") {
                                continue;
                        }
                        else if (tag == "NumPts") {
                                in >> m_numPts;
                        }
                        else if (tag == "DimSize") {
                                in >> m_dimSize;
                        }
                        else if (tag == "BucketSize") {
                                in >> m_bucketSize;
                        }
                        else if (tag == "BBoxLow") {
                                try {
                                        m_bboxLow.readFrom(in);
                                }
                                catch (FD::GeneralException *e) {
                                        throw e->add(new FD::GeneralException("In KDTree::readFrom : ",__FILE__,__LINE__));
                                }
                        }
                        else if (tag == "BBoxHigh") {
                                try {
                                        m_bboxHigh.readFrom(in);
                                }
                                catch (FD::GeneralException *e) {
                                        throw e->add(new FD::GeneralException("In KDTree::readFrom : ",__FILE__,__LINE__));
                                }
                        }
                        else if (tag == "KDRoot") {
                                // Allocate root node
                                try {
                                        m_root = new KDNode<KDData>();
                                        m_root->readFrom(in);
                                }
                                catch (FD::GeneralException *e) {
                                        throw e->add(new FD::GeneralException("In KDTree::readFrom : ",__FILE__,__LINE__));
                                }
                        }
                        else {
                                throw new FD::GeneralException ("KDTree::readFrom : Unknown argument: " + tag,__FILE__,__LINE__);
                        }

                        if (!in) {
                                throw new FD::GeneralException ("KDTree::readFrom : Parse error trying to build " + tag,__FILE__,__LINE__);
                        }

                        in >> tag;
                        if (tag != ">") {
                                throw new FD::GeneralException ("KDTree::readFrom : Parse error: '>' expected ",__FILE__,__LINE__);
                        }
                }
        }

        /*
                NNSearch

                Nearest neighbor search routine where the actual NN points are returned.
                This function requires the tree to be already built.

                PARAMETERS
                i_qp : query point given to find its nearest neighbor in the tree points
                                NOTE: the query point dimension size must be the same as the tree
                                                points.
                i_maxNumNodes : maximum number of leaf nodes to visit before to stop the
                                                search. If it is set to 0, than all the tree can be
                                                searched.
                i_errBnd : error bound to determine if we should visit a child node.
                                        A value of 0 means the search will visit all the potential
                                        nodes.
                i_numNN : number of nearest neighbors to look for
                i_squareDist : is the return NN distances in squared value?
                o_nnDist : output array to contain the NN distances.
                                        NOTE: we assume enough memeory has been allocated
                o_nn : output array to contain the NN.
                                NOTE: we assume enough memeory has been allocated
        */
        void NNSearch(KDPoint<KDData> *i_qp, int i_maxNumNodes, double i_errBnd,
                int i_numNN, bool i_squareDist, double *o_nnDist, KDPoint<KDData> *o_nn)
        {
                int numVisited = 0;
                double maxErr = (1.0 + i_errBnd)*(1.0 + i_errBnd);
                double dist =  BBoxDist2Point(i_qp);
                KDPQueue<double, KDPoint<KDData> > nnQueue(i_numNN);

                // Sanity check
                if (m_root == NULL) {
                        throw new FD::GeneralException ("KDTree::NNSearch : nearest neighbor search requires the KDTree to contain reference points. ",__FILE__,__LINE__);
                }

                if (i_qp == NULL) {
                        throw new FD::GeneralException ("KDTree::NNSearch : NULL query point. ",__FILE__,__LINE__);
                }

                if (o_nnDist == NULL) {
                        throw new FD::GeneralException ("KDTree::NNSearch : NULL nearest neighbor distances array. ",__FILE__,__LINE__);
                }

                if (o_nn == NULL) {
                        throw new FD::GeneralException ("KDTree::NNSearch : NULL nearest neighbor points array. ",__FILE__,__LINE__);
                }

                if (i_qp->GetDimSize() != m_dimSize) {
                        throw new FD::GeneralException ("KDTree::NNSearch : query point dimension size differs from KDTree reference points. ",__FILE__,__LINE__);
                }

                if (i_numNN == 0) {
                        return;
                }

                if (i_numNN > m_numPts) {
                        throw new FD::GeneralException ("KDTree::NNSearch : cannot search for more nearest neighbors than the number of reference points in KDTree. ",__FILE__,__LINE__);
                }

                // Start search at root node
                m_root->NNSearch(i_qp, i_maxNumNodes, numVisited, maxErr, dist, &nnQueue);

                // Copy NN and distances
                if (i_squareDist) {
                        for (int i=0; i<i_numNN; ++i) {
                                o_nnDist[i] = nnQueue.GetIthKey(i);
                                o_nn[i] = *(nnQueue.GetIthElement(i));
                        }
                }
                else {
                        for (int i=0; i<i_numNN; ++i) {
                                o_nnDist[i] = sqrt(nnQueue.GetIthKey(i));
                                o_nn[i] = *(nnQueue.GetIthElement(i));
                        }
                }
        }

        /*
                NNSearch

                Nearest neighbor search routine where the indexes to the NN reference points
                array used to build the tree are returned.
                This function requires the tree to be already built.

                PARAMETERS
                i_qp : query point given to find its nearest neighbor in the tree points
                                NOTE: the query point dimension size must be the same as the tree
                                                points.
                i_maxNumNodes : maximum number of leaf nodes to visit before to stop the
                                                search. If it is set to 0, than all the tree can be
                                                searched.
                i_errBnd : error bound to determine if we should visit a child node.
                                        A value of 0 means the search will visit all the potential
                                        nodes.
                i_numNN : number of nearest neighbors to look for
                i_squareDist : is the return NN distances in squared value?
                o_nnDist : output array to contain the NN distances.
                                        NOTE: we assume enough memeory has been allocated
                o_nnIdx : output array to contain the NN indexes.
                                NOTE: we assume enough memeory has been allocated
        */
        void NNSearch(KDPoint<KDData> *i_qp, int i_maxNumNodes, double i_errBnd,
                int i_numNN, bool i_squareDist, double *o_nnDist, int *o_nnIdx)
        {
                int numVisited = 0;
                double maxErr = (1.0 + i_errBnd)*(1.0 + i_errBnd);
                double dist =  BBoxDist2Point(i_qp);
                KDPQueue<double, int> nnQueue(i_numNN);

                // Sanity check
                if (m_root == NULL) {
                        throw new FD::GeneralException ("KDTree::NNSearch : nearest neighbor search requires the KDTree to contain reference points. ",__FILE__,__LINE__);
                }

                if (i_qp == NULL) {
                        throw new FD::GeneralException ("KDTree::NNSearch : NULL query point. ",__FILE__,__LINE__);
                }

                if (o_nnDist == NULL) {
                        throw new FD::GeneralException ("KDTree::NNSearch : NULL nearest neighbor distances array. ",__FILE__,__LINE__);
                        
                }

                if (o_nnIdx == NULL) {
                        throw new FD::GeneralException ("KDTree::NNSearch : NULL nearest neighbor indexes array. ",__FILE__,__LINE__);
                }

                if (i_qp->GetDimSize() != m_dimSize) {
                        throw new FD::GeneralException ("KDTree::NNSearch : query point dimension size differs from KDTree reference points. ",__FILE__,__LINE__);
                }

                if (i_numNN == 0) {
                        return;
                }

                if (i_numNN > m_numPts) {
                        throw new FD::GeneralException ("KDTree::NNSearch : cannot search for more nearest neighbors than the number of reference points in KDTree. ",__FILE__,__LINE__);
                }

                // Start search at root node
                m_root->NNSearch(i_qp, i_maxNumNodes, numVisited, maxErr, dist, &nnQueue);

                // Copy NN and distances
                if (i_squareDist) {
                        for (int i=0; i<i_numNN; ++i) {
                                o_nnDist[i] = nnQueue.GetIthKey(i);
                                o_nnIdx[i] = *(nnQueue.GetIthElement(i));
                        }
                }
                else {
                        for (int i=0; i<i_numNN; ++i) {
                                o_nnDist[i] = sqrt(nnQueue.GetIthKey(i));
                                o_nnIdx[i] = *(nnQueue.GetIthElement(i));
                        }
                }
        }

private:
        /*
                BBoxDist2Point

                Compute the distance from a given point to the current tree
                bounding box including all the reference points

                PARAMETERS
                i_qp : query point used to compute the distance

                RETURNS
                Distance from the given point to the bounding box
        */
        double BBoxDist2Point(KDPoint<KDData> *i_qp)
        {
                // Sum of squared distances
                double dist = 0.0;
                double tmp;

                for (int d = 0; d < m_dimSize; ++d) {
                        if (i_qp->GetCoord(d) < m_bboxLow.GetCoord(d)) {
                                // Query point is left of BBox
                                tmp = m_bboxLow.GetCoord(d) - i_qp->GetCoord(d);
                                dist += tmp*tmp;
                        }
                        else if (i_qp->GetCoord(d) > m_bboxHigh.GetCoord(d)) {
                                // Query point is right of BBox
                                tmp = i_qp->GetCoord(d) - m_bboxHigh.GetCoord(d);
                                dist += tmp*tmp;
                        }
                }

                return dist;
        }

        /*
                SetBBoxes

                Compute a bounding box including all the reference points.
                The bounding box is kept in member data.

                PARAMETERS
                i_points : given points used to build the tree structure
                i_ptsIdx : indexes to the points in the given array
        */
        void SetBBoxes(KDData **i_points, int *i_ptsIdx)
        {
                // find smallest enclosing rectangle
                for (int d=0; d<m_dimSize; ++d) {
                        // lower bound on dimension d
                        KDData low = i_points[i_ptsIdx[0]][d];
                        // upper bound on dimension d
                        KDData high = i_points[i_ptsIdx[0]][d];

                        for (int i=0; i<m_numPts; ++i) {
                                KDData cur = i_points[i_ptsIdx[i]][d];
                                if (cur < low) {
                                        low = cur;
                                }
                                else if (cur > high) {
                                        high = cur;
                                }
                        }

                        m_bboxLow.SetCoord(d, low);
                        m_bboxHigh.SetCoord(d, high);
                }
        }

        /*
                RecurBuildTree

                Recursive routine to build the tree. When the number of points
                is less or equal to the leaf node bucket size, we make a new
                leaf node. Otherwise, we build a new internal node. To construct
                an internal node, we need to find the cutting plane that will
                indicate which points go to the left side and which go to the
                right. The cutting plane is computed using Bentley's standard
                splitting routine for kd-trees.

                PARAMETERS
                i_points : given points used to build the tree structure
                i_ptsIdx : indexes to the points in the given array
                i_nPts : number of points left
        */
        KDNode<KDData> *RecurBuildTree(KDData **i_points, int *i_ptsIdx, int i_nPts)
        {
                // If i_nPts small, make a leaf node
                if (i_nPts <= m_bucketSize) {
                        if (i_nPts == 0) {
                                // Return a reference to an empty leaf node
                                return NULL;
                        }
                        else {
                                // Return a new leaf node
                                return new KDNode<KDData>(i_points, i_ptsIdx, i_nPts , m_dimSize);
                        }
                }

                // i_nPts large, make a splitting node
                // Cutting dimension
                int cd;
                // Cutting value
                KDData cv;
                // Number of data on left side of the cut
                int lowNum;
                // Low and high children
                KDNode<KDData> *lowNode, *highNode;

                // Find splitting node
                StdSplit(i_points, i_ptsIdx, i_nPts, cd, cv, lowNum);

                // Save bounds for cutting dimension
                KDData lv = m_bboxLow.GetCoord(cd);
                KDData hv = m_bboxHigh.GetCoord(cd);

                // Build left (low) subtree
                m_bboxHigh.SetCoord(cd, cv);
                lowNode = RecurBuildTree(i_points, i_ptsIdx, lowNum);

                // Restore bounds
                m_bboxHigh.SetCoord(cd, hv);

                // Build right (high) subtree
                m_bboxLow.SetCoord(cd, cv);
                highNode = RecurBuildTree(i_points, i_ptsIdx + lowNum, i_nPts-lowNum);
                // Restore bounds
                m_bboxLow.SetCoord(cd, lv);

                // Create the internal node
                return new KDNode<KDData>(cd, cv, lv, hv, lowNode, highNode);
        }

        /*
                StdSplit

                Bentley's standard splitting routine for kd-trees. We find the dimension
                with the greatest spread, and split just before the median point along
                this dimension.

                PARAMETERS
                i_points : given points used to build the tree structure
                i_ptsIdx : indexes to the points in the given array
                i_nPts : number of points in given array
                o_cd : cutting dimension (returned)
                o_cv : cutting value (returned)
                o_lowNum : number of points on low (left) side
        */
        void StdSplit(KDData **i_points, int *i_ptsIdx, int i_nPts, int &o_cd,
                KDData &o_cv, int &o_lowNum)
        {
                // Find dimension with maximum spread
                o_cd = MaxSpread(i_points, i_ptsIdx, i_nPts);

                // Median rank
                o_lowNum = i_nPts >> 1;

                MedianSplit(i_points, i_ptsIdx, i_nPts, o_cd, o_cv, o_lowNum);
        }

        /*
                MaxSpread

                Computes the dimension with the greatest spread in a given
                points array.

                PARAMETERS
                i_points : given points used to build the tree structure
                i_ptsIdx : indexes to the points in the given array
                i_nPts : number of points in given array

                RETURNS
                Dimension with the greatest spread
        */
        int MaxSpread(KDData **i_points, int *i_ptsIdx, int i_nPts)
        {
                int maxDim = 0;
                KDData maxSprd = (KDData)0;

                if (i_nPts == 0) {
                        return maxDim;
                }

                for (int d=0; d<m_dimSize; ++d) {
                        KDData sprd = DimSpread(i_points, i_ptsIdx, i_nPts, d);

                        if (sprd > maxSprd) {
                                maxSprd = sprd;
                                maxDim = d;
                        }
                }

                return maxDim;
        }

        /*
                DimSpread

                Computes the total spread along a given dimension in a given
                points array.

                PARAMETERS
                i_points : given points used to build the tree structure
                i_ptsIdx : indexes to the points in the given array
                i_nPts : number of points in given array
                i_dim : given dimension where to compute total spread

                RETURNS
                Total spread along given dimension
        */
        KDData DimSpread(KDData **i_points, int *i_ptsIdx, int i_nPts, int i_dim)
        {
                KDData min = i_points[i_ptsIdx[0]][i_dim];
                KDData max = i_points[i_ptsIdx[0]][i_dim];

                for (int i=0; i<i_nPts; ++i) {
                        KDData cur = i_points[i_ptsIdx[i]][i_dim];

                        if (cur < min) {
                                min = cur;
                        }
                        else if (cur > max) {
                                max = cur;
                        }
                }

                // Total spread over the dimension is the extremum difference
                return (max - min);
        }

        /*
                MedianSplit

                Split points array about its median into a subarray of points about
                an element of given rank in dimension i_dim. Here the points indexes
                are swap accordingly and the cutting value is computed.

                PARAMETERS
                i_points : given points used to build the tree structure
                i_ptsIdx : indexes to the points in the given array (modified)
                i_nPts : number of points in given array
                i_dim : given dimension to consider
                o_cv : cutting value (returned)
                i_lowIdx : number of points on low (left) side
        */
        void MedianSplit(KDData **i_points, int *i_ptsIdx, int i_nPts, int i_dim,
                KDData &o_cv, int i_lowIdx)
        {
                // left end of current data
                int l = 0;
                // right end of current data
                int r = i_nPts-1;

                while (l < r) {
                        // select middle as pivot
                        int i = (r+l)/2;
                        int k = 0;

                        // Pivot needs to be less than the last value
                        if (i_points[i_ptsIdx[i]][i_dim] > i_points[i_ptsIdx[r]][i_dim]) {
                                SwapPtsIdx(i_ptsIdx,i,r);
                        }

                        // move pivot to first position
                        SwapPtsIdx(i_ptsIdx,l,i);

                        KDData pivot = i_points[i_ptsIdx[l]][i_dim];
                        i = l;
                        k = r;

                        // Pivot around its value
                        while (1) {
                                while (i_points[i_ptsIdx[++i]][i_dim] < pivot);
                                while (i_points[i_ptsIdx[--k]][i_dim] > pivot);
                                if (i < k) {
                                        SwapPtsIdx(i_ptsIdx,i,k);
                                }
                                else {
                                        break;
                                }
                        }

                        // Place pivot at k
                        SwapPtsIdx(i_ptsIdx,l,k);

                        // Try to get the median
                        if (k > i_lowIdx) {
                                r = k-1;
                        }
                        else if (k < i_lowIdx) {
                                l = k+1;
                        }
                        else {
                                // Got it
                                break;
                        }
                }

                // Search for next smaller data
                if (i_lowIdx > 0) {
                        KDData c = i_points[i_ptsIdx[0]][i_dim];
                        int k = 0;

                        for (int i = 1; i < i_lowIdx; i++) {
                                if (i_points[i_ptsIdx[i]][i_dim] > c) {
                                        c = i_points[i_ptsIdx[i]][i_dim];
                                        k = i;
                                }
                        }

                        SwapPtsIdx(i_ptsIdx, i_lowIdx-1, k);
                }

                // Cut value is midpoint value
                o_cv = (i_points[i_ptsIdx[i_lowIdx-1]][i_dim] + i_points[i_ptsIdx[i_lowIdx]][i_dim])/(KDData)2;
        }

        /*
                SwapPtsIdx

                Swap element i_i with i_j in given array.

                PARAMETERS
                i_ptsIdx : array used to swap (modified)
                i_i : index of first element in array
                i_j : index of second element in array
        */
        void SwapPtsIdx(int *i_ptsIdx, int i_i, int i_j)
        {
                int tmp = i_ptsIdx[i_i];
                i_ptsIdx[i_i] = i_ptsIdx[i_j];
                i_ptsIdx[i_j] = tmp;
        }

private:
        // Number of reference points in tree structure
        int m_numPts;
        // Dimension size of reference points
        int m_dimSize;
        // Bucket size of leaf nodes
        int m_bucketSize;
        // Low corner point of bounding box including all the
        // reference points
        KDPoint<KDData> m_bboxLow;
        // High corner point of bounding box including all the
        // reference points
        KDPoint<KDData> m_bboxHigh;
        // Reference to root node
        KDNode<KDData> *m_root;

        // NN Search data member
        // Number of leaf nodes so far visisted
        int m_numVisited;
        // Maximum error to consider the exploration of the next node
        double m_maxErr;
};

}//namespace RobotFlow

#endif

---