更新 function.py

function.py

@@ -1,6 +1,6 @@
 import numpy as np
 import os
-
+from scipy.spatial import KDTree
 """
 实验说明:
 - 自行编写实现数据加载与 kNN 的手写数字分类算法。
@@ -41,17 +41,11 @@ def classify(input_vector, dataSet, labels, k=2):
     Returns:
     - int: 预测的分类标签。
     """
-    dataSet_size = dataSet.shape[0]
-    # 计算欧氏距离
-    diff_mat = np.tile(input_vector, (dataSet_size, 1)) - dataSet
-    sq_diff_mat = diff_mat ** 2
-    sq_distances = sq_diff_mat.sum(axis=1)
-    distances = np.sqrt(sq_distances)
-    # 对距离进行排序
-    sorted_indices = distances.argsort()
+    tree = KDTree(dataSet)
+    distances, indices = tree.query(input_vector, k=k)
     class_count = {}
     for i in range(k):
-        vote_label = labels[sorted_indices[i]]
+        vote_label = labels[indices[i]]
         class_count[vote_label] = class_count.get(vote_label, 0) + 1
     sorted_class_count = sorted(class_count.items(), key=lambda x: x[1], reverse=True)
     return sorted_class_count[0][0]