include/shark/Algorithms/NearestNeighbors/NearestNeighbors.h

Go to the documentation of this file.

//===========================================================================
//===========================================================================

#ifndef SHARK_ALGORITHMS_NEARESTNEIGHBORS_NEARESTNEIGHBORS_H
#define SHARK_ALGORITHMS_NEARESTNEIGHBORS_NEARESTNEIGHBORS_H


#include <boost/intrusive/rbtree.hpp>
#include <shark/Models/Kernels/AbstractKernelFunction.h>
#include <shark/Models/Trees/KDTree.h>
#include <shark/Models/Trees/LCTree.h>
#include <shark/Models/Trees/KHCTree.h>


namespace shark {


template <class InputT>
class IterativeNNQuery
{
    typedef BinaryTree<InputT> TreeType;
    typedef AbstractKernelFunction<InputT> KernelType;

public:
    IterativeNNQuery(const TreeType* tree, InputT const& point, const KernelType* kernel = NULL)
    : m_reference(point)
    , m_nextIndex(0)
    , mp_trace(NULL)
    , mep_head(NULL)
    , m_squaredRadius(0.0)
    , m_neighbors(0)
    , m_kernel(kernel)
    {
        // Initialize the recursion trace: descend to the
        // leaf covering the reference point and queue it.
        // The parent of this leaf becomes the "head".
        mp_trace = new TraceNode(tree, NULL, m_reference);
        TraceNode* tn = mp_trace;
        while (tree->hasChildren())
        {
            if (tree->left()->hasChildren())
                tn->mep_left = new TraceNode(tree->left(), tn, m_reference);
            else
                tn->mep_left = new TraceLeaf(tree->left(), m_kernel, tn, m_reference);
            if (tree->right()->hasChildren())
                tn->mep_right = new TraceNode(tree->right(), tn, m_reference);
            else
                tn->mep_right = new TraceLeaf(tree->right(), m_kernel, tn, m_reference);
            bool left = tree->isLeft(m_reference);
            tn = (left ? tn->mep_left : tn->mep_right);
            tree = (left ? tree->left() : tree->right());
        }
        mep_head = tn->mep_parent;
        insertIntoQueue((TraceLeaf*)tn);
        m_squaredRadius = mp_trace->squaredRadius(m_reference);
    }

    ~IterativeNNQuery()
    { delete mp_trace; }


    std::size_t neighbors() const
    { return m_neighbors; }

    std::size_t next(double& dist, const InputT*& point)
    {
        if (m_neighbors >= mp_trace->m_tree->size()) throw SHARKEXCEPTION("[IterativeNNQuery::next] no more neighbors available");

        assert(! m_queue.empty());

        // Check whether the current node has points
        // left, or whether it should be discarded.
        if (m_neighbors > 0)
        {
            TraceLeaf& q = *m_queue.begin();
            if (m_nextIndex < q.m_tree->size())
            {
                m_neighbors++;
                std::size_t index = q.m_tree->index(m_nextIndex);
                dist = std::sqrt(q.m_squaredPtDistance);
                point = &q.m_tree->point(m_nextIndex);
                m_nextIndex++;
                return index;
            }
            else m_queue.erase(q);
        }

        if (m_queue.empty() || (*m_queue.begin()).m_squaredPtDistance > m_squaredRadius)
        {
            // enqueue more points
            TraceNode* tn = mep_head;
            while (tn != NULL)
            {
                enqueue(tn);
                if (tn->m_status == COMPLETE) mep_head = tn->mep_parent;
                tn = tn->mep_parent;
            }

            // re-compute the radius
            m_squaredRadius = mp_trace->squaredRadius(m_reference);
        }

        m_neighbors++;
        TraceLeaf& q = *m_queue.begin();
        std::size_t index = q.m_tree->index(0);
        dist = std::sqrt(q.m_squaredPtDistance);
        point = &q.m_tree->point(0);
        m_nextIndex = 1;
        return index;
    }

    std::size_t queuesize() const
    { return m_queue.size(); }

private:
    enum Status
    {
        NONE,            // no points of this node have been queued yet
        PARTIAL,         // some of the points of this node have been queued
        COMPLETE,        // all points of this node have been queued
    };

    class TraceNode
    {
    public:
        TraceNode(const TreeType* tree, TraceNode* parent, const InputT& ref)
        : m_tree(tree)
        , m_status(NONE)
        , mep_parent(parent)
        , mep_left(NULL)
        , mep_right(NULL)
        , m_squaredDistance(tree->squaredDistanceLowerBound(ref))
        { }

        ~TraceNode()
        {
            if (mep_left != NULL) delete mep_left;
            if (mep_right != NULL) delete mep_right;
        }


        inline double squaredRadius(const InputT& ref) const
        {
            if (m_status == NONE) return m_squaredDistance;
            else if (m_status == PARTIAL)
            {
                double l = mep_left->squaredRadius(ref);
                double r = mep_right->squaredRadius(ref);
                return std::min(l, r);
            }
            else return 1e100;
        }

        const TreeType* m_tree;

        Status m_status;

        TraceNode* mep_parent;

        TraceNode* mep_left;

        TraceNode* mep_right;

        double m_squaredDistance;
    };

    typedef boost::intrusive::set_base_hook<> HookType;

    class TraceLeaf : public TraceNode, public HookType
    {
    public:
        TraceLeaf(const TreeType* tree, const KernelType* kernel, TraceNode* parent, const InputT& ref)
        : TraceNode(tree, parent, ref)
        , m_squaredPtDistance(kernel != NULL ? kernel->featureDistanceSqr((*tree)(0), ref) : shark::distanceSqr((*tree)(0), ref))
        { }

        ~TraceLeaf() { }


        inline bool operator < (TraceLeaf const& rhs) const
        {
            if (m_squaredPtDistance == rhs.m_squaredPtDistance) return (this->m_tree < rhs.m_tree);
            return (m_squaredPtDistance < rhs.m_squaredPtDistance);
        }

        double m_squaredPtDistance;
    };

    void insertIntoQueue(TraceLeaf* leaf)
    {
        m_queue.insert_unique(*leaf);

        // traverse up the tree, updating the state
        TraceNode* tn = leaf;
        tn->m_status = COMPLETE;
        while (true)
        {
            TraceNode* par = tn->mep_parent;
            if (par == NULL) break;
            if (par->m_status == NONE)
            {
                par->m_status = PARTIAL;
                break;
            }
            else if (par->m_status == PARTIAL)
            {
                if (par->mep_left == tn)
                {
                    if (par->mep_right->m_status == COMPLETE) par->m_status = COMPLETE;
                    else break;
                }
                else
                {
                    if (par->mep_left->m_status == COMPLETE) par->m_status = COMPLETE;
                    else break;
                }
            }
            tn = par;
        }
    }

    void enqueue(TraceNode* tn)
    {
        // check whether this node needs to be enqueued
        if (tn->m_status == COMPLETE) return;
        if (! m_queue.empty() && tn->m_squaredDistance >= (*m_queue.begin()).m_squaredPtDistance) return;

        const TreeType* tree = tn->m_tree;
        if (tree->hasChildren())
        {
            // extend the tree at need
            if (tn->mep_left == NULL)
            {
                if (tree->left()->hasChildren())
                    tn->mep_left = new TraceNode(tree->left(), tn, m_reference);
                else
                    tn->mep_left = new TraceLeaf(tree->left(), m_kernel, tn, m_reference);
            }
            if (tn->mep_right == NULL)
            {
                if (tree->right()->hasChildren())
                    tn->mep_right = new TraceNode(tree->right(), tn, m_reference);
                else
                    tn->mep_right = new TraceLeaf(tree->right(), m_kernel, tn, m_reference);
            }

            // first descend into the closer sub-tree
            if (tree->isLeft(m_reference))
            {
                // left first
                enqueue(tn->mep_left);
                enqueue(tn->mep_right);
            }
            else
            {
                // right first
                enqueue(tn->mep_right);
                enqueue(tn->mep_left);
            }
        }
        else
        {
            TraceLeaf* leaf = (TraceLeaf*)tn;
            insertIntoQueue(leaf);
        }
    }

    typedef boost::intrusive::rbtree<TraceLeaf> QueueType;

    InputT m_reference;

    QueueType m_queue;

    std::size_t m_nextIndex;

    TraceNode* mp_trace;

    TraceNode* mep_head;

    double m_squaredRadius;

    std::size_t m_neighbors;

    const KernelType* m_kernel;
};


}
#endif