diff --git a/function.py b/function.py
index 8cdca64..82b0c00 100644
--- a/function.py
+++ b/function.py
@@ -59,15 +59,21 @@ def classify(input_vector, dataSet, labels, k=2):
     Returns:
     - int: 预测的分类标签。
     """
+    if dataSet.ndim > 2:
+        dataSet = dataSet.reshape(dataSet.shape[0], -1)
+
     mean = np.mean(dataSet, axis=0)
     covmat = np.cov(dataSet, rowvar=0)
     inv_covmat = np.linalg.inv(covmat)
     
+    # 马氏距离计算
     dists = [mahalanobis_distance(input_vector, mean, inv_covmat)]
-    
     for data in dataSet:
         dists.append(mahalanobis_distance(data, mean, inv_covmat))
     
+    # 移除第一个距离，因为它是输入向量的马氏距离
+    dists.pop(0)
+    
     k_indices = np.argsort(dists)[:k]
     k_nearest_labels = [labels[i] for i in k_indices]
     
@@ -79,7 +85,7 @@ def classify(input_vector, dataSet, labels, k=2):
     if most_common:
         return most_common[0][0]
     else:
-        return None 
+        return None
 
 
 # 这里也写自己的代码实现