diff --git a/function.py b/function.py
index a9b6f26..edc5b20 100644
--- a/function.py
+++ b/function.py
@@ -39,7 +39,7 @@ name = "向兴宇"
 
 
 # 这里写自己的代码实现
-def classify0(inX, dataSet, labels, k=8, _cache={"VI": None}):
+def classify0(inX, dataSet, labels, k=8, _cache={"VI": None, "dataSet": None}):
     """
     KNN 分类算法。
 
@@ -61,14 +61,13 @@ def classify0(inX, dataSet, labels, k=8, _cache={"VI": None}):
     diff = dataSet - inX
 
     # 使用矩阵运算计算马氏距离
-    # distances = np.einsum("ij,jk,ik->i", diff, _cache["VI"], diff)
+    distances = np.einsum("ij,jk,ik->i", diff, _cache["VI"], diff)
+
+    # 对距离排序，获取前 k 个最近的样本索引
+    sortedDistIndices = np.argpartition(distances, 8)[:8]
 
     # 统计前 k 个样本的标签
-    classCount = np.bincount(
-        labels[
-            np.argpartition(np.einsum("ij,jk,ik->i", diff, _cache["VI"], diff), 8)[:8]
-        ]
-    )
+    classCount = np.bincount(labels[sortedDistIndices])
 
     # 返回出现次数最多的标签
     return np.argmax(classCount)