diff --git a/function.py b/function.py
index f0ed13d..5ea580b 100644
--- a/function.py
+++ b/function.py
@@ -60,10 +60,10 @@ def classify0(inX, dataSet, labels, k=8):
     distances = np.sqrt(np.sum(diff @ VI * diff, axis=1))
 
     # 对距离排序，获取前 k 个最近的样本索引
-    sortedDistIndices = np.argsort(distances)
+    sortedDistIndices = np.argsort(distances)[:k]
 
     # 统计前 k 个样本的标签
-    classCount = np.bincount(labels[sortedDistIndices[:k]])
+    classCount = np.bincount(labels[sortedDistIndices], minlength=np.max(labels) + 1)
 
     # 返回出现次数最多的标签
     return np.argmax(classCount)