如何为基于TF-IDF、PCA与K-Means的文本聚类散点图添加图例

嘿，我来帮你搞定这个图例的问题！你现在的散点图已经按聚类结果正确上色了，但缺个能说明“哪种颜色对应哪类文档”的图例，其实只要稍微调整下绘图代码就能轻松解决～

核心思路 #

原来你用plt.scatter一次性绘制所有点的方式，没办法为每个聚类单独生成图例项。所以我们需要按聚类分组绘制，给每个聚类的点指定对应的文档类型作为标签，这样就能自动生成清晰的图例啦。

具体实现步骤 #

1. 先明确聚类与文档类型的对应关系 #

首先你需要确定每个聚类ID（0-3）对应的文档类型，比如你已经知道：

• 聚类0 → 文档类型1
• 聚类1 → 文档类型2
• 聚类2 → 文档类型3
• 聚类3 → 文档类型4

把这个对应关系写成字典：

# 替换成你实际的文档类型名称
cluster_to_type = {
    0: "文档类型1",
    1: "文档类型2",
    2: "文档类型3",
    3: "文档类型4"
}

2. 重写绘图代码，分组绘制并添加图例 #

替换你原来的plt.scatter代码，改成按聚类循环绘制：

import matplotlib.pyplot as plt

# 可选：给每个聚类指定固定颜色，方便对应记忆
colors = ["green", "red", "blue", "orange"]

# 循环每个聚类ID，单独绘制该聚类的点
for cluster_id in range(4):
    # 筛选出当前聚类的所有点
    cluster_mask = (y_means == cluster_id)
    # 绘制散点，指定颜色和对应文档类型作为label
    plt.scatter(
        data2D[cluster_mask, 0], 
        data2D[cluster_mask, 1],
        c=colors[cluster_id],
        label=cluster_to_type[cluster_id],
        zorder=0
    )

# 添加图例，还可以给图例加标题
plt.legend(title="文档类型")

# 可选：添加图表标题和坐标轴标签
plt.title("文本聚类PCA散点图")
plt.xlabel("PCA 第一成分")
plt.ylabel("PCA 第二成分")

plt.show()

3. 如果不知道聚类对应哪种文档类型？ #

要是你还没明确每个聚类对应的文档类型，可以先统计每个聚类中占比最高的文档类型（前提是你有每个文档的类型标签列表doc_types）：

import pandas as pd

# 把数据整理成DataFrame方便统计
df = pd.DataFrame({
    "pca1": data2D[:, 0],
    "pca2": data2D[:, 1],
    "cluster": y_means,
    "doc_type": doc_types  # 你的文档类型列表，要和documents一一对应
})

# 自动统计每个聚类中占比最高的文档类型
cluster_type_map = df.groupby("cluster")["doc_type"].agg(
    lambda x: x.value_counts().index[0]
).to_dict()

# 之后用cluster_type_map替换前面的cluster_to_type即可

这样修改后，你的散点图就会带有清晰的图例，每个颜色对应的文档类型一目了然啦～

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