diff --git a/function.py b/function.py
index 21e1b8f..ec67f48 100644
--- a/function.py
+++ b/function.py
@@ -1,5 +1,6 @@
 import numpy as np
 import os
+from scipy.spatial import KDTree
 
 """
 实验说明:
@@ -41,11 +42,9 @@ def classify(input_vector, dataSet, labels, k=5):
     Returns:
     - int: 预测的分类标签。
     """
-    num_samples = dataSet.shape[0]
-    diffs = np.tile(input_vector, (num_samples, 1)) - dataSet
-    distances = np.sum(diffs**2, axis=1)**0.5
-    sorted_indices = np.argsort(distances)
-    k_nearest_labels = [labels[i] for i in sorted_indices[:k]]
+    kdtree = KDTree(dataSet)
+    distances, indices = kdtree.query(input_vector, k=k)
+    k_nearest_labels = [labels[i] for i in indices]
     unique_labels, counts = np.unique(k_nearest_labels, return_counts=True)
     return unique_labels[np.argmax(counts)]
 
@@ -87,49 +86,35 @@ def load_dataset(directory):
             labels.append(label)
     return np.array(data), np.array(labels)
 
+def load_test_dataset(directory):
+    return load_dataset(directory)
+
+def calculate_accuracy():
+    train_data, train_labels = load_dataset("trainingDigits")
+    test_data, test_labels = load_test_dataset("testDigits")
+    correct_count = 0
+    for i in range(len(test_data)):
+        predicted_label = classify(test_data[i].reshape(1, 1024), train_data, train_labels)
+        if predicted_label == test_labels[i]:
+            correct_count += 1
+    accuracy = correct_count / len(test_data)
+    print(f"分类精度: {accuracy}")
+
+def test_different_k():
+    train_data, train_labels = load_dataset("trainingDigits")
+    test_data, test_labels = load_test_dataset("testDigits")
+    k_values = [1, 3, 5, 7, 9]
+    for k in k_values:
+        correct_count = 0
+        for i in range(len(test_data)):
+            predicted_label = classify(test_data[i].reshape(1, 1024), train_data, train_labels, k=k)
+            if predicted_label == test_labels[i]:
+                correct_count += 1
+        accuracy = correct_count / len(test_data)
+        print(f"k = {k}时的分类精度: {accuracy}")
 
 if __name__ == "__main__":
-    # 以下代码可用于测试，可自行随意修改
-    train_data, train_labels = load_dataset("trainingDigits")
-
-    test_data = (
-        "00000000000000110000000000000000"
-        "00000000000011111000000000000000"
-        "00000000000011111100000000000000"
-        "00000000001111111110000000000000"
-        "00000000001111111110000000000000"
-        "00000000001111111100000000000000"
-        "00000000011111111000000000000000"
-        "00000000111111111000000000000000"
-        "00000000111111110000000000000000"
-        "00000000111111110000000000000000"
-        "00000000111111110000000000000000"
-        "00000000111111110000000000000000"
-        "00000000111111100000000000000000"
-        "00000001111111100000000000000000"
-        "00000001111111100000000000000000"
-        "00000001111111100000000000000000"
-        "00000001111111111111000000000000"
-        "00000001111111111111111000000000"
-        "00000001111111111111111100000000"
-        "00000011111111111111111110000000"
-        "00000011111111111111111110000000"
-        "00000001111111111111111110000000"
-        "00000001111111111111111110000000"
-        "00000000111111111100111111100000"
-        "00000000111111111000111111100000"
-        "00000000011111111001111111100000"
-        "00000000011111111111111111000000"
-        "00000000001111111111111111000000"
-        "00000000000111111111111111000000"
-        "00000000000011111111111100000000"
-        "00000000000001111111110000000000"
-        "00000000000000000111100000000000"
-    )
-    predicted_label = classify(
-        np.array([int(char) for char in test_data]).reshape(1, 1024),
-        train_data,
-        train_labels,
-    )
-
-    print(f"预测的标签: {predicted_label}")
\ No newline at end of file
+    # 测试不同k值的分类精度
+    test_different_k()
+    # 计算分类精度
+    calculate_accuracy()
\ No newline at end of file