diff --git a/function.py b/function.py
index 80ff2f8..c72d396 100644
--- a/function.py
+++ b/function.py
@@ -29,49 +29,61 @@ id = "202210423215"
 name = "章程鸿·"
 
 
-"""
-参数说明:
-- input_vector (numpy array): 输入待分类的样本。
-- dataSet (numpy array): 训练数据集。
-- labels (list or numpy array): 训练集对应的标签。
-- k (int): kNN 算法参数，请**自行设置默认值**，评测时仅会使用其默认值，不修改默认值将**无法通过**评测，**成绩为零**。
-返回值:
-- int: 预测的分类标签。
-"""
+import numpy as np
+import os
+from collections import Counter
 
 
-# 这里写自己的代码实现
-import numpy as np  
-import os  
-from collections import Counter  
+def classify(input_vector, dataSet, labels, k=2):
+    """
+    使用 KNN 算法对输入样本进行分类。
 
-def classify(input_vector, dataSet, labels, k=3):  
-    """使用 KNN 算法对输入样本进行分类。"""
-    # 计算协方差矩阵和其逆矩阵
-    cov_matrix = np.cov(dataSet, rowvar=False)
-    inv_cov_matrix = np.linalg.inv(cov_matrix + 1e-5 * np.eye(cov_matrix.shape[0]))  # 添加小常数以防止奇异矩阵
+    参数:
+            - input_vector (numpy array): 输入待分类的样本（1xN）。
+            - dataSet (numpy array): 训练数据集（MxN）。
+    - labels (list or numpy array): 训练集对应的标签。
+    - k (int): kNN 算法参数，默认值为 2。
 
+    返回:
+    - int: 预测的分类标签。
+    """
     # 计算马氏距离
-    diffs = input_vector - dataSet
-    distances = np.sqrt(np.einsum('ij,jk->i', diffs, np.dot(inv_cov_matrix, diffs.T)))
+    centered_dataSet = dataSet - np.mean(dataSet, axis=0)
+    centered_input_vector = input_vector - np.mean(dataSet, axis=0)
+    cov_matrix = np.cov(centered_dataSet.T)
+    if np.linalg.det(cov_matrix) == 0:
+        raise ValueError("协方差矩阵不可逆")
+    inv_cov_matrix = np.linalg.inv(cov_matrix)
+    mahalanobis_distances = np.sqrt(np.sum(centered_input_vector.dot(inv_cov_matrix) * centered_input_vector, axis=1))
 
-    sorted_indices = distances.argsort()
+    sorted_indices = mahalanobis_distances.argsort()
+    k_nearest_distances = mahalanobis_distances[sorted_indices[:k]]
     k_nearest_labels = labels[sorted_indices[:k]]
-    
-    # 计算权重（距离的倒数），避免除以零
-    weights = 1 / (distances[sorted_indices[:k]] + 1e-5)
+
+    # 计算权重（距离的倒数）
+    weights = 1 / k_nearest_distances
     # 归一化权重
     weights /= weights.sum()
-    
+
     # 计算加权投票
     weighted_votes = np.zeros((10,))
     for label, weight in zip(k_nearest_labels, weights):
         weighted_votes[label] += weight
-    
+
     return np.argmax(weighted_votes)
 
+
 def load_dataset(directory):
-    """加载数据集中的图像，并将其转换为特征向量和标签。"""
+    """
+    加载数据集中的图像，并将其转换为特征向量和标签。
+
+    参数:
+    - directory (str): 数据文件夹路径。
+
+    返回:
+            - np.array: 特征数据（Mx1024）。
+            - np.array: 标签数据（Mx1）。
+    """
     feature_list = []
     label_list = []
     for file in os.listdir(directory):
@@ -79,24 +91,32 @@ 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
 
-def img_to_vector(filename):  
-    """将二值图像转换为向量。"""
-    vector = np.zeros((1024,))  # 创建一个1024维的零向量
+
+def img_to_vector(filename):
+    """
+    将二值图像转换为向量。
+
+    Args:
+    - filename (str): 图像文件路径。
+
+    Returns:
+    - np.array: 二值图像向量。
+    """
+    vector = np.zeros((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()])
+            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 = (
@@ -133,9 +153,11 @@ if __name__ == "__main__":
         "00000000000001111111110000000000"
         "00000000000000000111100000000000"
     )
+    predicted_label = classify(
+        np.array([int(char) for char in test_data]).reshape(1, 1024),
+        train_data,
+        train_labels,
+    )
 
-    # 对测试数据进行归一化
-    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}")
+