Library version updated to v1.0.0.

The new leaf_ranges method of the kd_tree returns the indices per leaf.
Renamed apply_one_space to apply_dim_space.

Author: Jaybro

CMakeLists.txt

@@ -6,7 +6,7 @@ include(${CMAKE_CURRENT_SOURCE_DIR}/cmake/utils.cmake)
 
 project(pico_tree
     LANGUAGES CXX
-    VERSION 0.8.3
+    VERSION 1.0.0
     DESCRIPTION "PicoTree is a C++ header only library for fast nearest neighbor searches and range searches using a KdTree."
     HOMEPAGE_URL "https://github.com/Jaybro/pico_tree")
 

examples/kd_tree/kd_tree_custom_metric.cpp

@@ -5,13 +5,13 @@
 // This example shows how to create a custom metric for the kd_tree. The kd_tree
 // does not support metrics that "normalize" the distance between two points for
 // performance reasons. This means that, for example, the L2 metric is not
-// supported, but there is support for the L1 or L2^2 metric.
+// supported, but that there is support for the L1 or L2^2 metric.
 //
 // There are two different categories of metrics: Euclidean and topological. A
 // Euclidean metric supports any R^n space where all axes are orthogonal with
 // respect to each other. A topological space is the same, but it allows
-// wrapping of coordinate values along any of the axes. An implementation for
-// both is provided below.
+// wrapping of coordinate values along any of the axes. An implementation of a
+// custom metric is provided for both below.
 
 // The LP^P metric is a generalization of the L2^2 metric.
 template <std::size_t P_>
@@ -72,7 +72,7 @@ struct metric_t2_squared {
   }
 
   template <typename UnaryPredicate_>
-  void apply_one_space([[maybe_unused]] int dim, UnaryPredicate_ p) const {
+  void apply_dim_space([[maybe_unused]] int dim, UnaryPredicate_ p) const {
     p(pico_tree::one_space_s1{});
   }
 };

setup.py

@@ -5,7 +5,7 @@
 setup(
     name='pico_tree',
     # The same as the CMake project version.
-    version='0.8.3',
+    version='1.0.0',
     description='PicoTree Python Bindings',
     author='Jonathan Broere',
     url='https://github.com/Jaybro/pico_tree',

src/pico_tree/pico_tree/internal/box.hpp

@@ -357,7 +357,7 @@ class metric_box_map {
     auto bd = [&](auto one_space) {
       contains = make_segment(min, max, one_space).contains(v);
     };
-    metric_.apply_one_space(dim, bd);
+    metric_.apply_dim_space(dim, bd);
     return contains;
   }
 

src/pico_tree/pico_tree/internal/kd_tree_data.hpp

@@ -1,5 +1,8 @@
 #pragma once
 
+#include <iterator>
+#include <vector>
+
 #include "pico_tree/core.hpp"
 #include "pico_tree/internal/memory.hpp"
 #include "pico_tree/internal/stream_wrapper.hpp"
@@ -9,13 +12,33 @@ namespace pico_tree::internal {
 //! \brief The data structure that represents a kd_tree.
 template <typename Node_, size_t Dim_>
 class kd_tree_data {
+  template <typename Iterator_>
+  class iterator_range {
+   public:
+    using iterator_type = Iterator_;
+    using difference_type =
+        typename std::iterator_traits<Iterator_>::difference_type;
+
+    constexpr iterator_range(Iterator_ begin, Iterator_ end)
+        : begin_(begin), end_(end) {}
+
+    constexpr Iterator_ begin() const { return begin_; }
+    constexpr Iterator_ end() const { return end_; }
+
+   private:
+    Iterator_ begin_;
+    Iterator_ end_;
+  };
+
  public:
   using index_type = typename Node_::index_type;
   using scalar_type = typename Node_::scalar_type;
   static size_t constexpr dim = Dim_;
   using box_type = internal::box<scalar_type, dim>;
   using node_type = Node_;
   using node_allocator_type = chunk_allocator<node_type, 256>;
+  using leaf_range_type =
+      iterator_range<typename std::vector<index_type>::const_iterator>;
 
   static kd_tree_data load(stream_wrapper& stream) {
     typename box_type::size_type sdim;
@@ -34,6 +57,12 @@ class kd_tree_data {
     data.write(stream);
   }
 
+  inline std::vector<leaf_range_type> leaf_ranges() const {
+    std::vector<leaf_range_type> ranges;
+    insert_leaf_range(root_node, ranges);
+    return ranges;
+  }
+
   //! \brief Sorted indices that refer to points inside points_.
   std::vector<index_type> indices;
   //! \brief Bounding box of the root node.
@@ -44,6 +73,19 @@ class kd_tree_data {
   node_type* root_node;
 
  private:
+  inline void insert_leaf_range(
+      node_type const* const node, std::vector<leaf_range_type>& ranges) const {
+    if (node->is_leaf() && !node->data.leaf.empty()) {
+      using difference_type = typename leaf_range_type::difference_type;
+      ranges.push_back(leaf_range_type(
+          indices.begin() + difference_type(node->data.leaf.begin_idx),
+          indices.begin() + difference_type(node->data.leaf.end_idx)));
+    } else {
+      insert_leaf_range(node->left, ranges);
+      insert_leaf_range(node->right, ranges);
+    }
+  }
+
   //! \brief Recursively reads the Node and its descendants.
   inline node_type* read_node(stream_wrapper& stream) {
     node_type* node = allocator.allocate();

src/pico_tree/pico_tree/internal/kd_tree_node.hpp

@@ -29,6 +29,9 @@ struct kd_tree_node_base {
 //! \brief Tree leaf data.
 template <typename Index_>
 struct kd_tree_leaf {
+  //! \brief Returns true if the leaf is empty.
+  constexpr bool empty() const { return begin_idx == end_idx; }
+
   //! \brief Begin of an index range.
   Index_ begin_idx;
   //! \brief End of an index range.

src/pico_tree/pico_tree/internal/kd_tree_search.hpp

@@ -206,7 +206,7 @@ class search_nearest_topological {
       scalar_type min, scalar_type max, scalar_type v, int dim) const {
     scalar_type d;
     auto bd = [&](auto one_space) { d = box_distance(min, max, v, one_space); };
-    metric_.apply_one_space(dim, bd);
+    metric_.apply_dim_space(dim, bd);
     return metric_(d);
   }
 

src/pico_tree/pico_tree/kd_tree.hpp

@@ -50,6 +50,8 @@ class kd_tree {
   using metric_type = Metric_;
   //! \brief Neighbor type of various search results.
   using neighbor_type = neighbor<index_type, scalar_type>;
+  //! \brief Leaf index range type.
+  using leaf_range_type = typename kd_tree_data_type::leaf_range_type;
 
   //! \brief Creates a kd_tree given \p space and \p stop_condition.
   //! \details The kd_tree takes \p space by value. This allows it to take
@@ -269,7 +271,8 @@ class kd_tree {
   //! within radius \p radius and stores the results in output vector \p n.
   //! \see template <typename P_, typename RandomAccessIterator_> void
   //! search_radius(P_ const&, scalar_type, std::vector<neighbor_type>&, bool)
-  //! const \see template <typename P_, typename RandomAccessIterator_> void
+  //! const
+  //! \see template <typename P_, typename RandomAccessIterator_> void
   //! search_nn(P_ const&, scalar_type, neighbor_type&) const
   template <typename P_>
   inline void search_radius(
@@ -314,6 +317,15 @@ class kd_tree {
         idxs)(data_.root_node);
   }
 
+  //! \brief Returns the index range for all non-empty leaves in the kd_tree.
+  //! \details The leaf order follows from a depth-first traversal of the
+  //! kd_tree. If the input bounds is a hyper cube, then the leaf order is
+  //! identical to that of the N-order curve (a rotated version of the Z-order
+  //! curve that splits nodes in reversed order).
+  inline std::vector<leaf_range_type> leaf_ranges() const {
+    return data_.leaf_ranges();
+  }
+
   //! \brief Point set used by the tree.
   inline space_type const& space() const { return space_; }
 

src/pico_tree/pico_tree/metric.hpp

@@ -214,7 +214,7 @@ struct metric_so2 {
   }
 
   template <typename UnaryPredicate_>
-  void apply_one_space([[maybe_unused]] int dim, UnaryPredicate_ p) const {
+  void apply_dim_space([[maybe_unused]] int dim, UnaryPredicate_ p) const {
     p(one_space_s1{});
   }
 };
@@ -247,7 +247,7 @@ struct metric_se2_squared {
   }
 
   template <typename UnaryPredicate_>
-  void apply_one_space(int dim, UnaryPredicate_ p) const {
+  void apply_dim_space(int dim, UnaryPredicate_ p) const {
     if (dim < 2) {
       p(one_space_r1{});
     } else {

test/pico_tree/kd_tree_test.cpp

@@ -144,3 +144,35 @@ TEST(KdTreeTest, WriteRead) {
 
   EXPECT_TRUE(std::filesystem::remove(filename));
 }
+
+TEST(KdTreeTest, LeafRanges) {
+  using point_type = pico_tree::point_2f;
+  using scalar_type = typename point_type::scalar_type;
+  using index_type = int;
+  std::size_t point_count = 100;
+  scalar_type area_size = 2;
+  std::vector<point_type> random =
+      pico_tree::generate_random_n<point_type>(point_count, area_size);
+
+  kd_tree<point_type> tree(random, pico_tree::max_leaf_depth_t(2));
+
+  auto leaf_ranges = tree.leaf_ranges();
+  EXPECT_EQ(leaf_ranges.size(), 4);
+  if (!leaf_ranges.empty()) {
+    std::ptrdiff_t sum_range_point_count = 0;
+    index_type index_sum = 0;
+    for (auto const& r : leaf_ranges) {
+      sum_range_point_count += std::distance(r.begin(), r.end());
+      for (int i : r) {
+        index_sum += i;
+      }
+    }
+    EXPECT_EQ(static_cast<std::size_t>(sum_range_point_count), point_count);
+    EXPECT_EQ(
+        static_cast<std::size_t>(index_sum),
+        (point_count - 1) * point_count / 2);
+    EXPECT_EQ(
+        std::distance(leaf_ranges.front().begin(), leaf_ranges.back().end()),
+        point_count);
+  }
+}