From 98a12515fec3c79021c8a56acab6354e0764c43a Mon Sep 17 00:00:00 2001
From: MaximSmolskiy <mithridatus@mail.ru>
Date: Thu, 16 Oct 2025 22:58:55 +0300
Subject: [PATCH] Fix ml::KDTree::findNearest

---
 modules/ml/src/knearest.cpp       | 17 +++++++++++++++--
 modules/ml/test/test_knearest.cpp |  7 ++++++-
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/modules/ml/src/knearest.cpp b/modules/ml/src/knearest.cpp
index 3d8f9b5d2e..bd9137b24b 100644
--- a/modules/ml/src/knearest.cpp
+++ b/modules/ml/src/knearest.cpp
@@ -396,8 +396,21 @@ public:
         {
             Mat _res, _nr, _d;
             tr.findNearest(test_samples.row(i), k, Emax, _res, _nr, _d, noArray());
-            res.push_back(_res.t());
-            _results.assign(res);
+            if( _results.needed() )
+            {
+                res.push_back(_res.t());
+                _results.assign(res);
+            }
+            if( _neighborResponses.needed() )
+            {
+                nr.push_back(_nr.t());
+                _neighborResponses.assign(nr);
+            }
+            if( _dists.needed() )
+            {
+                d.push_back(_d.t());
+                _dists.assign(d);
+            }
         }
 
         return result; // currently always 0
diff --git a/modules/ml/test/test_knearest.cpp b/modules/ml/test/test_knearest.cpp
index 80baed9626..56d3e94269 100644
--- a/modules/ml/test/test_knearest.cpp
+++ b/modules/ml/test/test_knearest.cpp
@@ -39,11 +39,16 @@ TEST(ML_KNearest, accuracy)
     {
         SCOPED_TRACE("KDTree");
         Mat neighborIndexes;
+        Mat neighborResponses;
+        Mat dists;
         float err = 1000;
         Ptr<KNearest> knn = KNearest::create();
         knn->setAlgorithmType(KNearest::KDTREE);
         knn->train(trainData, ml::ROW_SAMPLE, trainLabels);
-        knn->findNearest(testData, 4, neighborIndexes);
+        knn->findNearest(testData, 4, neighborIndexes, neighborResponses, dists);
+        EXPECT_EQ(neighborIndexes.size(), Size(4, pointsCount));
+        EXPECT_EQ(neighborResponses.size(), Size(4, pointsCount * 2));
+        EXPECT_EQ(dists.size(), Size(4, pointsCount));
         Mat bestLabels;
         // The output of the KDTree are the neighbor indexes, not actual class labels
         // so we need to do some extra work to get actual predictions