如何基于Numpy等基础库为KMeans代码实现Dunn指数与惯性计算

需求说明 #

现有基于Numpy、Matplotlib和CSV库实现的KMeans聚类代码，需添加惯性（Inertia）和Dunn指数的计算功能，且只能使用这三个库完成。

修改后的完整代码 #

import numpy as np
import matplotlib.pyplot as plt
import csv

def load_data(file_path):
    data = []
    with open(file_path, 'r') as csvfile:
        csvreader = csv.reader(csvfile)
        next(csvreader)
        for row in csvreader:
            data.append([float(row[0]), float(row[1])])
    return np.array(data)

def calculate_distances(data, centers):
    data_with_distances = data.copy()
    num_centers = centers.shape[0]
    
    for i in range(num_centers):
        distances = np.sqrt(((data - centers[i]) ** 2).sum(axis=1))
        data_with_distances = np.column_stack((data_with_distances, distances))
    
    return data_with_distances

def get_clusters(data_with_distances):
    num_clusters = data_with_distances.shape[1] - 2 
    cluster_masks = []
    for i in range(num_clusters):
        mask = data_with_distances[:, i+2] == np.min(data_with_distances[:, 2:num_clusters+2], axis=1)
        cluster_masks.append(mask)
    
    clusters = [data_with_distances[mask, :] for mask in cluster_masks]
    return clusters

def calculate_centers(clusters):
    centers = np.array([cluster.mean(axis=0)[:2] for cluster in clusters])
    return centers

def plot_clusters(clusters, centers):
    colors = ['blue', 'red', 'green', 'orange', 'pink', 'gray', 'brown', 'cyan', 'magenta', 'yellow']
    for i, cluster in enumerate(clusters):
        plt.scatter(cluster[:, 0], cluster[:, 1], color=colors[i])
    for center in centers:
        plt.scatter(center[0], center[1], color='purple', marker='*', s=150)
    plt.xlabel('Household Total Assets')
    plt.ylabel('Annual Household Income')
    plt.title(f'K-means Clustering (k={len(clusters)})')
    plt.show()

def calculate_inertia(clusters, centers):
    """计算惯性：所有样本到其所属簇中心的距离平方和"""
    inertia = 0.0
    for i, cluster in enumerate(clusters):
        points = cluster[:, :2]  # 提取原始数据列
        center = centers[i]
        # 计算每个点到中心的距离平方并累加
        distances_sq = ((points - center) ** 2).sum(axis=1)
        inertia += distances_sq.sum()
    return inertia

def calculate_dunn_index(clusters):
    """计算Dunn指数：最小簇间距离 / 最大簇内直径，值越大聚类效果越好"""
    num_clusters = len(clusters)
    if num_clusters < 2:
        return 0.0
    
    # 计算每个簇的直径（簇内任意两点间的最大距离）
    cluster_diameters = []
    for cluster in clusters:
        points = cluster[:, :2]
        # 生成簇内所有点对的距离矩阵
        dist_matrix = np.sqrt(((points[:, np.newaxis] - points) ** 2).sum(axis=2))
        cluster_diameters.append(dist_matrix.max())
    max_diameter = max(cluster_diameters)
    
    # 计算所有簇对之间的最小距离
    min_inter_dist = float('inf')
    for i in range(num_clusters):
        for j in range(i+1, num_clusters):
            points_i = clusters[i][:, :2]
            points_j = clusters[j][:, :2]
            # 计算两个簇所有点对的距离矩阵，取最小值
            cross_dist_matrix = np.sqrt(((points_i[:, np.newaxis] - points_j) ** 2).sum(axis=2))
            current_min = cross_dist_matrix.min()
            if current_min < min_inter_dist:
                min_inter_dist = current_min
    
    # 处理特殊情况：若有簇只有一个点（直径为0）
    if max_diameter == 0:
        return float('inf') if min_inter_dist > 0 else 0.0
    
    return min_inter_dist / max_diameter

def run(data, num_clusters, max_iterations=100):
    current_centers = np.random.permutation(data)[:num_clusters]
    
    for iteration in range(max_iterations):
        data_with_distances = calculate_distances(data, current_centers)
        clusters = get_clusters(data_with_distances)
        new_centers = calculate_centers(clusters)
        # 提前终止：若中心不再变化
        if np.allclose(current_centers, new_centers):
            break
        current_centers = new_centers

    plot_clusters(clusters, current_centers)
    # 计算并输出指标
    inertia = calculate_inertia(clusters, current_centers)
    dunn_idx = calculate_dunn_index(clusters)
    print(f"=== 聚类数k={num_clusters} ===")
    print(f"惯性值：{inertia:.2f}")
    print(f"Dunn指数：{dunn_idx:.4f}\n")
    
def main(file_path):
    data = load_data(file_path)
    for num_clusters in range(2, 11):
        run(data, num_clusters)

file_path = 'assessment2dmv.csv'
main(file_path)

关键功能说明 #

惯性（Inertia）

• 定义：所有样本到其所属簇中心的距离平方和，值越小说明簇内样本越紧凑。
• 实现逻辑：遍历每个簇，计算簇内所有点到中心的距离平方，累加得到总和。

Dunn指数

• 定义：最小簇间距离除以最大簇内直径，值越大说明簇间区分度越高、簇内越紧凑，聚类效果越好。
• 实现逻辑：
  1. 计算每个簇的直径：簇内任意两点间的最大距离；
  2. 计算所有簇对之间的最小距离：不同簇的任意两点间的最小距离；
  3. 用最小簇间距离除以最大簇直径得到指数值。

内容的提问来源于stack exchange，提问作者cosmic coder