diff --git a/function.py b/function.py
index 82b0c00..68030b7 100644
--- a/function.py
+++ b/function.py
@@ -46,24 +46,17 @@ def mahalanobis_distance(x, mean, inv_covmat):
     m_dist = np.sqrt(np.dot(np.dot(delta, inv_covmat), delta.T))
     return m_dist
 
-def classify(input_vector, dataSet, labels, k=2):
-    """
-    KNN 分类算法。
-
-    Args:
-    - input_vector (numpy array): 输入待分类的样本。
-    - dataSet (numpy array): 训练数据集。
-    - labels (list or numpy array): 训练集对应的标签。
-    - k (int): kNN 算法参数，请**自行设置默认值**，评测时仅会使用其默认值，不修改默认值将**无法通过**评测，**成绩为零**。
-
-    Returns:
-    - int: 预测的分类标签。
-    """
+def classify(input_vector, dataSet, labels, k=2, reg=1e-5):
+    # 确保dataSet是二维的，每个样本是一维的
     if dataSet.ndim > 2:
         dataSet = dataSet.reshape(dataSet.shape[0], -1)
-
+    
     mean = np.mean(dataSet, axis=0)
     covmat = np.cov(dataSet, rowvar=0)
+    
+    # 正则化协方差矩阵以避免奇异矩阵
+    covmat += np.eye(covmat.shape[0]) * reg
+    
     inv_covmat = np.linalg.inv(covmat)
     
     # 马氏距离计算