KdTree.h

Go to the documentation of this file.

/*************************************************************************
> File Name: KdTree.h
> Project Name: CubbyFlow
> Author: Chan-Ho Chris Ohk
> Purpose: Generic k-d tree structure.
> Created Time: 2017/12/05
> Copyright (c) 2018, Chan-Ho Chris Ohk
*************************************************************************/
#ifndef CUBBYFLOW_KDTREE_H
#define CUBBYFLOW_KDTREE_H

#include <Core/BoundingBox/BoundingBox.h>
#include <Core/Vector/Vector.h>

namespace CubbyFlow
{
    template <typename T, size_t K>
    class KdTree final
    {
    public:
        using Point = Vector<T, K>;
        using BBox = BoundingBox<T, K>;

        struct Node
        {
            size_t flags = 0;

            size_t child = std::numeric_limits<size_t>::max();

            size_t item = std::numeric_limits<size_t>::max();

            Point point;

            Node();

            void InitLeaf(size_t it, const Point& pt);

            void InitInternal(size_t axis, size_t it, size_t c, const Point& pt);

            bool IsLeaf() const;
        };

        using ContainerType = std::vector<Point>;
        using Iterator = typename ContainerType::iterator;
        using ConstIterator = typename ContainerType::const_iterator;

        using NodeContainerType = std::vector<Node>;
        using NodeIterator = typename NodeContainerType::iterator;
        using ConstNodeIterator = typename NodeContainerType::const_iterator;

        KdTree();

        void Build(const ConstArrayAccessor1<Point>& points);

        void ForEachNearbyPoint(
            const Point& origin, T radius,
            const std::function<void(size_t, const Point&)>& callback) const;

        bool HasNearbyPoint(const Point& origin, T radius) const;

        size_t GetNearestPoint(const Point& origin) const;

        Iterator begin();

        Iterator end();

        ConstIterator begin() const;

        ConstIterator end() const;

        NodeIterator BeginNode();

        NodeIterator EndNode();

        ConstNodeIterator BeginNode() const;

        ConstNodeIterator EndNode() const;

        void Reserve(size_t numPoints, size_t numNodes);

    private:
        std::vector<Point> m_points;
        std::vector<Node> m_nodes;

        size_t Build(size_t nodeIndex, size_t* itemIndices, size_t nItems, size_t currentDepth);
    };
}

#include <Core/Geometry/KdTree-Impl.h>

#endif