更新 function.py

function.py

@@ -5,7 +5,9 @@ import numpy as np
 import os
 from collections import Counter
 from sklearn.preprocessing import MinMaxScaler
+from sklearn.neighbors import KDTree
 
+# 加载数据集
 def load_dataset(directory):
     file_list = os.listdir(directory)
     data_matrix = []
@@ -13,15 +15,14 @@ def load_dataset(directory):
 
     for file_name in file_list:
         file_path = os.path.join(directory, file_name)
-        # 读取txt文件并将其转化为向量
         vector = img_to_vector(file_path)
         data_matrix.append(vector)
-        # 假设文件名格式为：{数字}_某些字符.txt，提取数字作为标签
         label = int(file_name.split('_')[0])
         labels.append(label)
 
     return np.array(data_matrix), np.array(labels)
 
+# 将32x32的二值图像转换为1x1024的向量
 def img_to_vector(file_path):
     vector = np.zeros((1, 1024))
     with open(file_path, 'r') as f:
@@ -31,26 +32,20 @@ def img_to_vector(file_path):
                 vector[0, 32 * i + j] = int(lines[i][j])
     return vector
 
-def classify(input_vector, dataSet, labels, k=7):
-    # 计算欧氏距离
-    distances = np.linalg.norm(dataSet - input_vector, axis=1)
-    # 获取距离最近的k个样本的标签
-    nearest_neighbors = np.argsort(distances)[:k]
-    nearest_labels = labels[nearest_neighbors]
-    
-    # 确保 nearest_labels 是一维的，转换为 Python 列表
-    nearest_labels = nearest_labels.flatten().tolist()
+# 使用 KDTree 进行 kNN 分类
+def classify(input_vector, tree, labels, k=3):
+    # 使用 KDTree 查找最近邻
+    distances, indices = tree.query(input_vector.reshape(1, -1), k=k)
+    nearest_labels = labels[indices.flatten()]
     
     # 投票，返回出现次数最多的标签
     most_common = Counter(nearest_labels).most_common(1)
     return most_common[0][0]
 
+# 数据归一化
 def normalize_data(dataSet):
-    # 将数据集 reshape 成二维数组
     num_samples = dataSet.shape[0]
     dataSet_reshaped = dataSet.reshape(num_samples, -1)
-    
-    # 对数据进行归一化处理
     scaler = MinMaxScaler()
     return scaler.fit_transform(dataSet_reshaped)
 
@@ -63,13 +58,16 @@ def main():
     training_data = normalize_data(training_data)
     test_data = normalize_data(test_data)
 
+    # 使用 KDTree 构建训练集的树结构
+    tree = KDTree(training_data)
+
     # 测试 kNN 算法
     total_tests = len(test_data)
     correct_predictions = 0
 
     for i in range(total_tests):
         test_sample = test_data[i]
-        prediction = classify(test_sample, training_data, training_labels, k=7)
+        prediction = classify(test_sample, tree, training_labels, k=3)
         actual_label = test_labels[i]
 
         print(f"测试样本 {i+1}: 预测值 = {prediction}, 实际值 = {actual_label}")