如何将K-Means/TF-IDF文本聚类的PCA可视化图按作者着色？

我首次使用K-Means、TF-IDF进行文档聚类，已完成文本文件聚类，并通过PCA降维可视化了聚类分布。现在希望把绘图的颜色标识从聚类类别改成文本对应的作者，而非聚类/主题，该怎么实现？

当前实现代码 #

file_list = glob.glob(os.path.join(os.getcwd(), "myFiles", "*.txt"))

dataset = []

for file_path in file_list:
    with open(file_path) as f_input:
        dataset.append(f_input.read())

vectorizer = TfidfVectorizer(stop_words='english')

vectorized_documents = vectorizer.fit_transform(dataset) 
  
pca = PCA(n_components=2) 
reduced_data = pca.fit_transform(vectorized_documents.toarray()) 
  
num_clusters = 7
kmeans = KMeans(n_clusters=num_clusters, n_init=5, 
                max_iter=500, random_state=42) 
kmeans.fit(vectorized_documents) 

# create a dataframe to store the results 
results = pd.DataFrame() 
results['document'] = dataset 
results['cluster'] = kmeans.labels_ 

# plot the results 
colors = ['black', 'red', 'green', 'yellow', 'blue', 'orange', 'purple'] 
cluster = ['0', '1','2', '3', '4', '5', '6'] 
for i in range(num_clusters): 
    plt.scatter(reduced_data[kmeans.labels_ == i, 0], 
                reduced_data[kmeans.labels_ == i, 1],  
                s=10, color=colors[i],  
                label=f' {cluster[i]}') 
plt.legend() 
plt.show()

效果对比 #

当前聚类着色效果图：

期望的作者着色效果图：

实现步骤 #

1. 提取作者信息 #

首先需要从文件名或文件内容中提取每个文档对应的作者。假设文件名包含作者信息（例如author1_doc1.txt），可通过解析文件名获取：

file_list = glob.glob(os.path.join(os.getcwd(), "myFiles", "*.txt"))
dataset = []
authors = []  # 新增存储作者的列表

for file_path in file_list:
    with open(file_path) as f_input:
        dataset.append(f_input.read())
    # 从文件名提取作者，根据实际文件名格式调整分割规则
    filename = os.path.basename(file_path)
    author = filename.split('_')[0]  # 示例：按下划线分割取第一部分作为作者
    authors.append(author)

如果作者信息在文件内容中，需修改逻辑从文本里解析对应字段。

2. 将作者信息加入结果DataFrame #

results = pd.DataFrame() 
results['document'] = dataset 
results['cluster'] = kmeans.labels_ 
results['author'] = authors  # 新增作者列

3. 按作者分组绘图 #

遍历所有唯一作者，为每个作者分配颜色并绘制对应的数据点：

# 获取所有唯一作者
unique_authors = results['author'].unique()
# 定义颜色列表（数量需≥作者数量，可按需扩展）
colors = ['black', 'red', 'green', 'yellow', 'blue', 'orange', 'purple', 'pink', 'gray']

# 按作者绘制散点图
for idx, author in enumerate(unique_authors):
    # 筛选当前作者对应的PCA降维数据
    author_data = reduced_data[results['author'] == author]
    plt.scatter(author_data[:, 0], author_data[:, 1],  
                s=10, color=colors[idx],  
                label=author) 

plt.legend(title='作者') 
plt.xlabel('PCA 维度1')
plt.ylabel('PCA 维度2')
plt.title('文档聚类可视化（按作者着色）')
plt.show()

优化方案：自动生成颜色 #

如果作者数量较多，可使用matplotlib的颜色映射自动生成颜色，无需手动定义颜色列表：

cmap = plt.cm.get_cmap('tab10', len(unique_authors))  # 选择颜色映射，数量匹配作者数

for idx, author in enumerate(unique_authors):
    author_data = reduced_data[results['author'] == author]
    plt.scatter(author_data[:, 0], author_data[:, 1],  
                s=10, color=cmap(idx),  
                label=author)

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