更新 function.py

function.py

@@ -43,7 +43,7 @@ name = "熊浩燃"
 
 
 # 这里写自己的代码实现
-def classify_optimized(input_vector, dataSet, labels, k=3, metric='euclidean', decision_rule='uniform'):
+def classify(input_vector, dataSet, labels, k=3, metric='euclidean', decision_rule='uniform', use_kdtree=False):
     """
     优化后的 KNN 分类算法，包括 k 值调优、距离度量、决策规则和加速方法。
     
@@ -54,11 +54,20 @@ def classify_optimized(input_vector, dataSet, labels, k=3, metric='euclidean', d
     - k (int): kNN 算法中的 k 值。
     - metric (str): 距离度量方式，默认是 'euclidean'，可以选择 'manhattan', 'cosine', 'mahalanobis' 等。
     - decision_rule (str): 决策规则，'uniform' 为普通投票，'weighted' 为加权投票。
+    - use_kdtree (bool): 是否使用 KDTree 加速。
     
     Returns:
     - int: 预测的分类标签。
     """
-    # 选择距离度量
+    # 如果使用 KDTree 加速
+    if use_kdtree:
+        tree = KDTree(dataSet)
+        distances, indices = tree.query(input_vector.reshape(1, -1), k=k)
+        nearest_labels = np.array(labels)[indices[0]]
+        predicted_label = np.bincount(nearest_labels).argmax()
+        return predicted_label
+
+    # 如果不使用 KDTree, 选择距离度量
     if metric == 'euclidean':
         distances = np.linalg.norm(dataSet - input_vector, axis=1)
     elif metric == 'manhattan':
@@ -82,27 +91,6 @@ def classify_optimized(input_vector, dataSet, labels, k=3, metric='euclidean', d
 
     return predicted_label
 
-# KDTree 加速 KNN 分类算法
-def classify_with_kdtree(input_vector, dataSet, labels, k=3):
-    """
-    使用 KDTree 加速 KNN 算法。
-    
-    Args:
-    - input_vector (numpy array): 输入待分类的样本。
-    - dataSet (numpy array): 训练数据集。
-    - labels (list or numpy array): 训练集对应的标签。
-    - k (int): kNN 算法中的 k 值。
-    
-    Returns:
-    - int: 预测的分类标签。
-    """
-    tree = KDTree(dataSet)
-    distances, indices = tree.query(input_vector.reshape(1, -1), k=k)
-    nearest_labels = np.array(labels)[indices[0]]
-    
-    predicted_label = np.bincount(nearest_labels).argmax()
-    return predicted_label
-
 # 这里也写自己的代码实现
 def load_dataset(directory):
     """
@@ -144,12 +132,11 @@ def img_to_vector(filename):
             vector[32 * i + j] = int(lines[i][j])
     return vector
 
-
 if __name__ == "__main__":
     # 测试代码
-    train_data, train_labels = load_dataset("trainingDigits")
+   train_data, train_labels = load_dataset("trainingDigits")
 
-    test_data = (
+   test_data = (
         "00000000000000110000000000000000"
         "00000000000011111000000000000000"
         "00000000000011111100000000000000"
@@ -187,9 +174,9 @@ if __name__ == "__main__":
     test_vector = np.array([int(char) for char in test_data]).reshape(1, 1024)
     
     # 使用优化后的 KNN 进行预测
-    predicted_label_optimized = classify_optimized(test_vector, train_data, train_labels, k=5, metric='cosine', decision_rule='weighted')
+    predicted_label_optimized = classify(test_vector, train_data, train_labels, k=5, metric='cosine', decision_rule='weighted', use_kdtree=False)
     print(f"优化后的预测标签: {predicted_label_optimized}")
     
     # 使用 KDTree 加速的 KNN 进行预测
-    predicted_label_kdtree = classify_with_kdtree(test_vector, train_data, train_labels, k=5)
+    predicted_label_kdtree = classify(test_vector, train_data, train_labels, k=5, use_kdtree=True)
     print(f"KDTree 加速后的预测标签: {predicted_label_kdtree}")