KdTree.hpp

Go to the documentation of this file.

// This code is based on Jet framework.
// Copyright (c) 2018 Doyub Kim
// CubbyFlow is voxel-based fluid simulation engine for computer games.
// Copyright (c) 2020 CubbyFlow Team
// Core Part: Chris Ohk, Junwoo Hwang, Jihong Sin, Seungwoo Yoo
// AI Part: Dongheon Cho, Minseo Kim
// We are making my contributions/submissions to this project solely in our
// personal capacity and are not conveying any rights to any intellectual
// property of any third parties.
 
#ifndef CUBBYFLOW_KDTREE_HPP
#define CUBBYFLOW_KDTREE_HPP
 
#include <Core/Array/ArrayView.hpp>
#include <Core/Geometry/BoundingBox.hpp>
#include <Core/Matrix/Matrix.hpp>
 
namespace CubbyFlow
{
template <typename T, size_t K>
class KdTree final
{
 public:
    using Point = Vector<T, K>;
    using BBox = BoundingBox<T, K>;
 
    struct Node
    {
        void InitLeaf(size_t it, const Point& pt);
 
        void InitInternal(size_t axis, size_t it, size_t c, const Point& pt);
 
        [[nodiscard]] bool IsLeaf() const;
 
        size_t flags = 0;
 
        size_t child = std::numeric_limits<size_t>::max();
 
        size_t item = std::numeric_limits<size_t>::max();
 
        Point point;
    };
 
    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;
 
    void Build(const ConstArrayView1<Point>& points);
 
    void ForEachNearbyPoint(
        const Point& origin, T radius,
        const std::function<void(size_t, const Point&)>& callback) const;
 
    [[nodiscard]] bool HasNearbyPoint(const Point& origin, T radius) const;
 
    [[nodiscard]] size_t GetNearestPoint(const Point& origin) const;
 
    [[nodiscard]] Iterator begin();
 
    [[nodiscard]] Iterator end();
 
    [[nodiscard]] ConstIterator begin() const;
 
    [[nodiscard]] ConstIterator end() const;
 
    [[nodiscard]] NodeIterator BeginNode();
 
    [[nodiscard]] NodeIterator EndNode();
 
    [[nodiscard]] ConstNodeIterator BeginNode() const;
 
    [[nodiscard]] ConstNodeIterator EndNode() const;
 
    void Reserve(size_t numPoints, size_t numNodes);
 
 private:
    [[nodiscard]] size_t Build(size_t nodeIndex, size_t* itemIndices,
                               size_t nItems, size_t currentDepth);
 
    std::vector<Point> m_points;
    std::vector<Node> m_nodes;
};
}  // namespace CubbyFlow
 
#include <Core/Geometry/KdTree-Impl.hpp>
 
#endif