当前位置：首页 > article >正文

自定义数据集，使用scikit-learn 中K均值包进行聚类

article 2025/2/3 11:44:40

代码：

# 导入必要的库
import matplotlib.pyplot as plt  # 用于绘制图形
from sklearn.cluster import KMeans  # KMeans 聚类算法
import numpy as np  # 数值计算库

# 定义 class1 到 class4 的数据点，模拟四个不同的类（每个类7个二维点）
class1_points = np.array([[1.9, 1.2],
                          [1.5, 2.1],
                          [1.9, 0.5],
                          [1.5, 0.9],
                          [0.9, 1.2],
                          [1.1, 1.7],
                          [1.4, 1.1]])

class2_points = np.array([[-1.9, 1.2],
                          [-1.5, 2.1],
                          [-1.9, 0.5],
                          [-1.5, 0.9],
                          [-0.9, 1.2],
                          [-1.1, 1.7],
                          [-1.4, 1.1]])

class3_points = np.array([[1.9, -1.2],
                          [1.5, -2.1],
                          [1.9, -0.5],
                          [1.5, -0.9],
                          [0.9, -1.2],
                          [1.1, -1.7],
                          [1.4, -1.1]])

class4_points = np.array([[-1.9, -1.2],
                          [-1.5, -2.1],
                          [-1.9, -0.5],
                          [-1.5, -0.9],
                          [-0.9, -1.2],
                          [-1.1, -1.7],
                          [-1.4, -1.1]])

# 将四个类的数据合并成一个大的数据集 'data'，这是所有数据点的集合
data = np.concatenate((class1_points, class2_points, class3_points, class4_points))

# 创建一个 1x2 的子图布局，figsize 控制图形大小
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))

# k 为聚类的数量，设定为 4，即我们预期将数据聚成四类
k = 4

# 在 ax1 上绘制所有数据点，s=50 控制点的大小
ax1.scatter(data[:, 0], data[:, 1], s=50)
ax1.set_title("Original Data")  # 设置图表标题

# 初始化 KMeans 聚类模型，设定聚类数量 k 和最大迭代次数为 30
km = KMeans(n_clusters=k, max_iter=30)

# 使用 KMeans 聚类算法对数据进行训练
km.fit(data)

# 获取聚类中心（每个聚类的质心）
centeroids = km.cluster_centers_

# 使用 KMeans 的 predict 方法预测每个数据点所属的簇
y_kmean = km.predict(data)

# 打印每个数据点所属的簇的索引
print(y_kmean)

# 在 ax2 上绘制聚类的结果
for i in range(k):
    # 提取当前簇的数据点
    cluster_points = data[y_kmean == i]
    # 获取当前簇的质心
    centroid = centeroids[i]
    # 将当前簇的每个点与质心连接，并用虚线绘制
    for cluster_point in cluster_points:
        ax2.plot([cluster_point[0], centroid[0]], [cluster_point[1], centroid[1]], 'k--')

# 在 ax2 上绘制所有数据点，并根据聚类结果着色
ax2.scatter(data[:, 0], data[:, 1], c=y_kmean, s=50)

# 在聚类结果图上绘制每个簇的质心（用黑色标记，点的大小为 100，透明度为 0.5）
ax2.scatter(centeroids[:, 0], centeroids[:, 1], c='black', s=100, alpha=0.5)

# 设置图表标题
ax2.set_title("KMeans Clustering Result")

# 显示绘制的图形
plt.show()

结果：