聚类算法评估指标之SSE代码实现

import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from sklearn.datasets import make_blobs
from sklearn.metrics import calinski_harabasz_score

# 1. 生成数据
print("=== 1. 生成数据 ===")
X, y_true = make_blobs(
    n_samples=500,
    centers=4,
    cluster_std=0.5,
    random_state=42
)
print(f"数据形状: {X.shape}")

# 2. SSE + 肘部法
print("\n=== 2. SSE + 肘部法 ===")

def demo_sse_elbow(X):
    """SSE + 肘部法选择最佳K值"""
    print("\n--- SSE + 肘部法 ---")
    
    # K值范围（SSE可以从1开始）
    k_range = range(1, 11)
    sse_list = []
    
    for k in k_range:
        # 创建KMeans对象
        kmeans = KMeans(n_clusters=k, random_state=42, n_init=10)
        # 训练并预测
        kmeans.fit(X)
        # 获取SSE值
        sse = kmeans.inertia_
        sse_list.append(sse)
        print(f"K={k}: SSE={sse:.4f}")
    
    # 绘制折线图
    plt.figure(figsize=(10, 6))
    plt.plot(k_range, sse_list, 'bo-')
    plt.xlabel('K值（簇数）')
    plt.ylabel('SSE（误差平方和）')
    plt.title('肘部法（Elbow Method）- SSE')
    plt.grid(True)
    plt.xticks(k_range)
    plt.savefig('sse_elbow.png', dpi=100)
    print("肘部法图已保存")
    
    # 找肘部点
    print("\n肘部点分析:")
    for i in range(1, len(sse_list)-1):
        change1 = sse_list[i-1] - sse_list[i]
        change2 = sse_list[i] - sse_list[i+1]
        if change1 > change2 * 2:
            print(f"  K={i+1} 可能是肘部点（下降变缓）")
    
    return sse_list

sse_list = demo_sse_elbow(X)

# 3. 手动计算SSE
print("\n=== 3. 手动计算SSE ===")

def calculate_sse_manual(X, labels, centroids):
    """手动计算SSE"""
    sse = 0
    k = len(centroids)
    
    for cluster_idx in range(k):
        cluster_points = X[labels == cluster_idx]
        centroid = centroids[cluster_idx]
        # 每个样本到质心的距离平方
        distances_sq = np.sum((cluster_points - centroid) ** 2, axis=1)
        # 簇内误差平方和
        cluster_sse = np.sum(distances_sq)
        sse += cluster_sse
    
    return sse

# 使用K=4
kmeans = KMeans(n_clusters=4, random_state=42, n_init=10)
labels = kmeans.fit_predict(X)
centroids = kmeans.cluster_centers_

sse_manual = calculate_sse_manual(X, labels, centroids)
print(f"手动计算SSE: {sse_manual:.4f}")
print(f"sklearn inertia_: {kmeans.inertia_:.4f}")

# 4. 不同数据集的SSE对比
print("\n=== 4. 不同数据集的SSE对比 ===")

# 紧凑数据
X_compact, _ = make_blobs(n_samples=500, centers=4, cluster_std=0.3, random_state=42)
# 松散数据
X_loose, _ = make_blobs(n_samples=500, centers=4, cluster_std=1.5, random_state=42)

kmeans_compact = KMeans(n_clusters=4, random_state=42, n_init=10)
kmeans_compact.fit(X_compact)

kmeans_loose = KMeans(n_clusters=4, random_state=42, n_init=10)
kmeans_loose.fit(X_loose)

print(f"紧凑数据SSE: {kmeans_compact.inertia_:.4f}")
print(f"松散数据SSE: {kmeans_loose.inertia_:.4f}")

# 5. 可视化对比
plt.figure(figsize=(12, 5))

plt.subplot(1, 2, 1)
plt.scatter(X_compact[:, 0], X_compact[:, 1], c='blue', alpha=0.6)
plt.title(f'紧凑数据 (SSE={kmeans_compact.inertia_:.2f})')

plt.subplot(1, 2, 2)
plt.scatter(X_loose[:, 0], X_loose[:, 1], c='red', alpha=0.6)
plt.title(f'松散数据 (SSE={kmeans_loose.inertia_:.2f})')

plt.tight_layout()
plt.savefig('sse_comparison.png', dpi=100)
print("SSE对比图已保存")

# 6. SSE公式说明
print("\n=== 6. SSE公式说明 ===")
print("""
SSE公式:

SSE = Σₖ₌₁ᴷ Σᵢ₌₁ⁿᵏ (pᵢ - μₖ)²

说明:
- K: 聚类中心的个数
- k: 第k个簇
- nₖ: 第k个簇的样本数
- pᵢ: 第i个样本
- μₖ: 第k个簇的质心

特点:
- 只考虑簇内距离（内聚）
- 没有考虑簇间距离（耦合）
- SSE越小越好
""")

# 7. 完整流程
def sse_evaluation_pipeline():
    """SSE评估完整流程"""
    print("=" * 50)
    print("SSE评估 - 完整流程")
    print("=" * 50)
    
    # 1. 生成数据
    print("\n1. 生成数据")
    X, _ = make_blobs(n_samples=500, centers=4, cluster_std=0.5, random_state=42)
    print(f"   数据形状: {X.shape}")
    
    # 2. 计算不同K值的SSE
    print("\n2. 计算不同K值的SSE")
    k_range = range(1, 11)
    sse_list = []
    
    for k in k_range:
        kmeans = KMeans(n_clusters=k, random_state=42, n_init=10)
        kmeans.fit(X)
        sse_list.append(kmeans.inertia_)
    
    # 3. 绘制肘部法图
    print("\n3. 绘制肘部法图")
    plt.figure(figsize=(10, 6))
    plt.plot(k_range, sse_list, 'bo-', linewidth=2, markersize=8)
    plt.xlabel('K值', fontsize=12)
    plt.ylabel('SSE', fontsize=12)
    plt.title('肘部法选择最佳K值', fontsize=14)
    plt.grid(True, alpha=0.3)
    plt.xticks(k_range)
    
    # 标记肘部点
    plt.annotate('肘部点 (K=4)', 
                 xy=(4, sse_list[3]), 
                 xytext=(6, sse_list[3] + 100),
                 arrowprops=dict(arrowstyle='->', color='red', lw=2),
                 fontsize=12, color='red')
    
    plt.savefig('sse_elbow_annotated.png', dpi=100)
    print("   肘部法图已保存")
    
    # 4. 选择最佳K值
    print("\n4. 选择最佳K值")
    print("   根据肘部法，K=4是最佳值")
    
    print("\n" + "=" * 50)
    print("评估完成!")
    print("=" * 50)

sse_evaluation_pipeline()