diff --git a/function.py b/function.py
index 83246af..5450c89 100644
--- a/function.py
+++ b/function.py
@@ -45,29 +45,15 @@ import numpy as np
 import os  
 from collections import Counter  
 
-  
-# 这里写自己的代码实现  
-def classify(input_vector, dataSet, labels, k=2):  
-    """  
-    使用 KNN 算法对输入样本进行分类。  
-  
-    参数:  
-	    - input_vector (numpy array): 输入待分类的样本（1xN）。  
-	    - dataSet (numpy array): 训练数据集（MxN）。  
-    - labels (list or numpy array): 训练集对应的标签。  
-    - k (int): kNN 算法参数，默认值为 2。  
-  
-    返回:  
-    - int: 预测的分类标签。  
-    """  
-    # 在 classify 函数中实现权重投票
+def classify(input_vector, dataSet, labels, k=3):  
+    """使用 KNN 算法对输入样本进行分类。"""
     distances = np.sum((input_vector - dataSet) ** 2, axis=1)
     sorted_indices = distances.argsort()
     k_nearest_distances = distances[sorted_indices[:k]]
     k_nearest_labels = labels[sorted_indices[:k]]
     
-    # 计算权重（距离的倒数）
-    weights = 1 / k_nearest_distances
+    # 计算权重（距离的倒数），避免除以零
+    weights = 1 / (k_nearest_distances + 1e-5)
     # 归一化权重
     weights /= weights.sum()
     
@@ -78,8 +64,8 @@ def classify(input_vector, dataSet, labels, k=2):
     
     return np.argmax(weighted_votes)
 
-# 在 load_dataset 函数中添加特征缩放
 def load_dataset(directory):
+    """加载数据集中的图像，并将其转换为特征向量和标签。"""
     feature_list = []
     label_list = []
     for file in os.listdir(directory):
@@ -87,74 +73,24 @@ def load_dataset(directory):
         feature_vector = img_to_vector(file_path)
         feature_list.append(feature_vector)
         label_list.append(int(file.split('_')[0]))
+    
     feature_matrix = np.array(feature_list)
     label_array = np.array(label_list)
-    
-    # 特征缩放
+
+    # 特征归一化
     feature_matrix = feature_matrix / 255.0
     return feature_matrix, label_array
-    # 计算输入向量与训练集之间的欧氏距离
-    distances = np.sum((input_vector - dataSet) ** 2, axis=1) ** 0.5
-    # 根据距离排序并获取索引
-    sorted_indices = distances.argsort()
-    # 选择距离最近的k个点
-    k_nearest_labels = labels[sorted_indices[:k]].tolist()  # 将 NumPy 数组转换为列表
-    # 统计k个点中各类别的数量
-    count = Counter(k_nearest_labels)
-    # 返回数量最多的类别标签
-    return count.most_common(1)[0][0]
 
-# 这里也写自己的代码实现  
-def load_dataset(directory):  
-    """  
-    加载数据集中的图像，并将其转换为特征向量和标签。  
-  
-    参数:  
-    - directory (str): 数据文件夹路径。  
-  
-    返回:  
-	    - np.array: 特征数据（Mx1024）。  
-	    - np.array: 标签数据（Mx1）。  
-    """  
-    # 初始化特征数据列表和标签列表
-    feature_list = []
-    label_list = []
-    # 遍历目录中的所有文件
-    for file in os.listdir(directory):
-        # 获取文件的完整路径
-        file_path = os.path.join(directory, file)
-        # 将图像文件转换为向量
-        feature_vector = img_to_vector(file_path)
-        # 将特征向量添加到列表中
-        feature_list.append(feature_vector)
-        # 获取文件名中的标签并添加到标签列表中
-        label_list.append(int(file.split('_')[0]))
-    # 将特征列表和标签列表转换为numpy数组
-    feature_matrix = np.array(feature_list)
-    label_array = np.array(label_list)
-    return feature_matrix, label_array
-# 这里还写自己的代码实现  
 def img_to_vector(filename):  
-    """  
-    将二值图像转换为向量。  
-  
-    Args:  
-    - filename (str): 图像文件路径。  
-  
-    Returns:  
-    - np.array: 二值图像向量。  
-    """  
-    # 初始化向量
-    vector = np.zeros((1024,))  # 创建一个一维数组
-    # 打开文件并逐行读取
+    """将二值图像转换为向量。"""
+    vector = np.zeros((1024,))  # 创建一个1024维的零向量
     with open(filename) as file:
         for i, line in enumerate(file):
-            # 将每行的数字字符转换为整数并存储到向量中
             vector[i*32:i*32+32] = np.array([int(char) for char in line.strip()])
     return vector
 
 if __name__ == "__main__":
-    # 以下代码可用于测试，可自行随意修改
+    # 以下代码可用于测试
     train_data, train_labels = load_dataset("trainingDigits")
 
     test_data = (
@@ -191,10 +127,9 @@ if __name__ == "__main__":
         "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
+    # 对测试数据进行归一化
+    input_vector = np.array([int(char) for char in test_data]).reshape(1, 1024) / 255.0
+
+    predicted_label = classify(input_vector, train_data, train_labels)
+    print(f"预测的标签: {predicted_label}")