From 38061ee738e16e555a64bb94d169530aece7c035 Mon Sep 17 00:00:00 2001
From: Gianluca Bigaglia <bigagliagianluca@gmail.com>
Date: Mon, 28 Apr 2025 14:51:48 +0200
Subject: [PATCH] Add a new API to retrieve both nearest indices and distances

---
 KDTree.cpp | 21 +++++++++++++++++++++
 KDTree.hpp | 10 ++++++++++
 2 files changed, 31 insertions(+)

diff --git a/KDTree.cpp b/KDTree.cpp
index 09adb92..ba02e50 100644
--- a/KDTree.cpp
+++ b/KDTree.cpp
@@ -221,6 +221,27 @@ pointIndexArr KDTree::nearest_pointIndices(point_t const& pt,
     return output;
 }
 
+std::pair<std::vector<double>, std::vector<size_t>> KDTree::nearest_indices_dists(
+    point_t const& pt, size_t const& num_nearest) {
+
+    size_t level = 0;
+    std::list<std::pair<KDNodePtr, double>> k_buffer{};
+    k_buffer.emplace_back(root_, dist2(static_cast<point_t>(*root_), pt));
+    knearest_(root_,       // beginning of tree
+              pt,          // point we are querying
+              level,       // start from level 0
+              num_nearest, // number of nearest neighbours to return in k_buffer
+              k_buffer);   // list of k nearest neigbours (to be filled)
+    std::vector<size_t> output_ids;
+    std::vector<double> output_dists;
+    std::for_each(k_buffer.begin(), k_buffer.end(),
+      [&output_ids, &output_dists](auto const& nodeptr_dist) {
+                    output_ids.push_back(nodeptr_dist.first->index);
+                    output_dists.push_back(nodeptr_dist.second);
+                  });
+    return {output_dists, output_ids};
+}
+
 pointVec KDTree::nearest_points(point_t const& pt, size_t const& num_nearest) {
     auto const k_nearest{nearest_pointIndices(pt, num_nearest)};
     pointVec k_nearest_points(k_nearest.size());
diff --git a/KDTree.hpp b/KDTree.hpp
index 8e587bb..cbc2486 100644
--- a/KDTree.hpp
+++ b/KDTree.hpp
@@ -106,6 +106,16 @@ class KDTree {
     pointIndexArr nearest_pointIndices(point_t const& pt,
                                        size_t const& num_nearest);
 
+    /// Gets index of points closest to the given input point and their
+    /// respective distance from it.
+    ///
+    /// @param pt input point.
+    /// @param num_nearest Number of nearest points to return.
+    /// @returns a vector containing the point distances and a vector
+    /// containing point indices.
+    std::pair<std::vector<double>, std::vector<size_t>>
+      nearest_indices_dists(point_t const& pt, size_t const& num_nearest);
+
     /// Get the nearest set of points to the given input point.
     ///
     /// @param pt input point.