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