From 1d8f5b6b0be0d02008e973027a6242da13c90e6b Mon Sep 17 00:00:00 2001
From: loulanyue <260355617@qq.com>
Date: Mon, 24 Aug 2026 15:40:39 +0800
Subject: [PATCH] fix(ml): add validation for matching dataSet and labels
 lengths in kNN

---
 src/algorithms/ml/knn/__test__/knn.test.js | 7 +++++++
 src/algorithms/ml/knn/kNN.js               | 4 ++++
 2 files changed, 11 insertions(+)

diff --git a/src/algorithms/ml/knn/__test__/knn.test.js b/src/algorithms/ml/knn/__test__/knn.test.js
index 1a7a117fe3..cb155e386a 100644
--- a/src/algorithms/ml/knn/__test__/knn.test.js
+++ b/src/algorithms/ml/knn/__test__/knn.test.js
@@ -28,6 +28,13 @@ describe('kNN', () => {
     expect(inconsistent).toThrow('Matrices have different shapes');
   });
 
+  it('should throw an error on mismatched dataSet and labels lengths', () => {
+    const mismatched = () => {
+      kNN([[1, 1], [2, 2]], [1], [1, 1]);
+    };
+    expect(mismatched).toThrow('The number of dataSet points must match the number of labels');
+  });
+
   it('should find the nearest neighbour', () => {
     let dataSet;
     let labels;
diff --git a/src/algorithms/ml/knn/kNN.js b/src/algorithms/ml/knn/kNN.js
index 350e7ce176..a21974ebc3 100644
--- a/src/algorithms/ml/knn/kNN.js
+++ b/src/algorithms/ml/knn/kNN.js
@@ -20,6 +20,10 @@ export default function kNN(
     throw new Error('Either dataSet or labels or toClassify were not set');
   }
 
+  if (dataSet.length !== labels.length) {
+    throw new Error('The number of dataSet points must match the number of labels');
+  }
+
   // Calculate distance from toClassify to each point for all dimensions in dataSet.
   // Store distance and point's label into distances list.
   const distances = [];