diff --git a/function.py b/function.py
index ce6cc91..0669ab1 100644
--- a/function.py
+++ b/function.py
@@ -59,7 +59,7 @@ def classify(input_vector, dataSet, labels, k=3):
     # 获取距离最近的 k 个样本的索引
     nearest_neighbors_indices = distances.argsort()[:k]
     # 获取这些样本的标签
-    nearest_labels = np.array(labels)[nearest_neighbors_indices]
+    nearest_labels = labels[nearest_neighbors_indices]
     # 统计出现次数最多的标签并返回
     predicted_label = np.bincount(nearest_labels).argmax()
     return predicted_label
@@ -108,8 +108,10 @@ def img_to_vector(filename):
 
 
 if __name__ == "__main__":
+    # 测试代码
     train_data, train_labels = load_dataset("trainingDigits")
-
+    
+    # 测试数据（缩减以减少运行时间）
     test_data = (
         "00000000000000110000000000000000"
         "00000000000011111000000000000000"