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 = [];