更新 function.py

function.py

@@ -1,5 +1,6 @@
 # function.py
 import numpy as np
+from collections import defaultdict
 
 """
 挑战说明:
@@ -21,7 +22,6 @@ import numpy as np
 - numpy: 2.1.1
 """
 
-
 """
 请在这里填写学号和姓名，不填写即成绩无效。
 """
@@ -39,66 +39,89 @@ name = "严锡好"
 - int: 预测的分类标签。
 """
 
-
 # 这里写自己的代码实现
-class KDNode:
-    def __init__(self, data, left=None, right=None, label=None):
-        self.data = data
-        self.left = left
-        self.right = right
-        self.label = label
+def mahalanobis_distance(x1, x2, cov_inv, mean):
+    """
+    计算两点之间的马氏距离。
+    
+    Args:
+    x1 (numpy array): 第一个点。
+    x2 (numpy array): 第二个点。
+    cov_inv (numpy array): 协方差矩阵的逆矩阵。
+    mean (numpy array): 数据集的均值。
+    
+    Returns:
+    float: 两点之间的马氏距离。
+    """
+    delta = x1 - x2
+    dist = np.dot(delta, np.dot(cov_inv, delta.T))  # 使用 dot 函数进行矩阵乘法
+    return np.sqrt(dist)
 
-class KDTree:
-    def __init__(self, data, labels, k=3):
-        self.k = k
-        self.root = self.build_tree(data, labels)
-
-    def build_tree(self, data, labels, depth=0):
-        if not data:
-            return None
-        indices = np.argsort(data[:, depth % data.shape[1]])
-        median_idx = indices[len(indices) // 2]
-        node = KDNode(data[median_idx], label=labels[median_idx])
-        node.left = self.build_tree(data[indices[:median_idx]], labels[indices[:median_idx]], depth + 1)
-        node.right = self.build_tree(data[indices[median_idx + 1:]], labels[indices[median_idx + 1:]], depth + 1)
-        return node
-
-    def find_nearest_neighbors(self, point, node=None, depth=0, neighbors=None):
-        if not node:
-            return
-        dist = np.linalg.norm(point - node.data)
-        if neighbors is None:
-            neighbors = [(dist, node.label)]
-        else:
-            neighbors = [(n[0], n[1]) for n in neighbors if n[0] < dist] + [(dist, node.label)]
-        neighbors.sort(key=lambda x: x[0])
-        if len(neighbors) < self.k:
-            self.find_nearest_neighbors(point, node.left, depth + 1, neighbors)
-            self.find_nearest_neighbors(point, node.right, depth + 1, neighbors)
-        else:
-            if depth % point.shape[0] != point.shape[0] - 1:
-                self.find_nearest_neighbors(point, node.left, depth + 1, neighbors)
-            if depth % point.shape[0] != point.shape[0] - 1:
-                self.find_nearest_neighbors(point, node.right, depth + 1, neighbors)
-
-    def predict(self, point):
-        neighbors = []
-        self.find_nearest_neighbors(point, self.root, 0, neighbors)
-        neighbors.sort(key=lambda x: x[0])
-        labels = [n[1] for n in neighbors]
-        return np.argmax(np.bincount(labels))
+def classify0(inX, dataSet, labels, k=3, cov_inv=None, mean=None):
+    """
+    使用马氏距离的 KNN 分类算法。
+    
+    Args:
+    inX (numpy array): 输入待分类的样本。
+    dataSet (numpy array): 训练数据集。
+    labels (list or numpy array): 训练集对应的标签。
+    k (int): kNN 算法参数，默认值为 3。
+    cov_inv (numpy array): 协方差矩阵的逆矩阵，如果为 None，则需要重新计算。
+    mean (numpy array): 数据集的均值，如果为 None，则需要重新计算。
+    
+    Returns:
+    int: 预测的分类标签。
+    """
+    if cov_inv is None or mean is None:  # 如果没有提供协方差逆矩阵或均值，则计算它们
+        cov_matrix = np.cov(dataSet, rowvar=False)
+        mean = np.mean(dataSet, axis=0)
+        cov_inv = np.linalg.inv(cov_matrix + np.eye(cov_matrix.shape[0]) * 1e-5)  # 加上一个小的单位矩阵以避免奇异矩阵
+    
+    # 使用向量化操作计算所有样本与待分类样本之间的马氏距离
+    diff = dataSet - mean  # 预先计算所有样本与均值的差
+    dist = np.sum(diff * cov_inv[None, :, :] * diff[:, None, :], axis=2)
+    
+    # 获取距离最近的 k 个样本的索引
+    sortedDistIndices = np.argsort(dist)[:, :k]
+    
+    # 使用 defaultdict 来简化类别计数逻辑
+    classCount = defaultdict(int)
+    for idx in sortedDistIndices:
+        classCount[labels[idx]] += 1
+    
+    # 选择出现次数最多的标签作为预测标签
+    sortedClassCount = sorted(classCount.items(), key=lambda x: x[1], reverse=True)
+    return sortedClassCount[0][0]
 
 def load_data(file_path):
+    """
+    加载 TSV 格式数据文件。
+    """
     data = np.loadtxt(file_path, delimiter="\t")
     features = data[:, :-1]
     labels = data[:, -1].astype(int)
     return features, labels
 
+def prepare_data(dataSet):
+    """
+    计算训练数据集的协方差矩阵的逆矩阵和均值。
+    
+    Args:
+    dataSet (numpy array): 训练数据集。
+    
+    Returns:
+    tuple: 包含均值和协方差矩阵的逆矩阵的元组。
+    """
+    cov_matrix = np.cov(dataSet, rowvar=False)
+    mean = np.mean(dataSet, axis=0)
+    cov_inv = np.linalg.inv(cov_matrix + np.eye(cov_matrix.shape[0]) * 1e-5)
+    return mean, cov_inv
+
 if __name__ == "__main__":
     train_file = "datingTrainSet2.tsv"
     train_data, train_labels = load_data(train_file)
-    kdt = KDTree(train_data, train_labels, k=3)
+    mean, cov_inv = prepare_data(train_data)  # 预先计算均值和协方差逆矩阵
 
     test_sample = np.array([65915, 3.868359, 0.128078])
-    predicted_label = kdt.predict(test_sample)
+    predicted_label = classify0(test_sample, train_data, train_labels, k=3, cov_inv=cov_inv, mean=mean)
     print(f"The predicted label for the test sample is: {predicted_label}")