Sklearn -按类别分组，并从每个数据帧类别中获得前n个单词？

from sklearn.feature_extraction.text import CountVectorizer

# 假设有一个包含多个文档的列表data，每个文档属于一个类别
data = ['文档1内容', '文档2内容', ...]

# 创建CountVectorizer对象
vectorizer = CountVectorizer()

# 将文本数据转换为词频矩阵
X = vectorizer.fit_transform(data)

from sklearn.feature_extraction.text import TfidfTransformer

# 创建TfidfTransformer对象
transformer = TfidfTransformer()

# 将词频矩阵转换为TF-IDF矩阵
X_tfidf = transformer.fit_transform(X)

import numpy as np

# 假设有一个包含类别标签的列表labels，与data中的文档一一对应
labels = ['类别1', '类别2', ...]

# 获取所有类别
categories = np.unique(labels)

# 遍历每个类别
for category in categories:
    # 获取该类别的文档索引
    indices = np.where(labels == category)[0]
    
    # 获取该类别的TF-IDF矩阵
    X_category = X_tfidf[indices]
    
    # 计算每个单词在该类别中的TF-IDF均值
    tfidf_mean = np.mean(X_category, axis=0)
    
    # 获取前n个单词的索引
    top_n_indices = np.argsort(tfidf_mean)[-n:]
    
    # 获取前n个单词
    top_n_words = np.array(vectorizer.get_feature_names())[top_n_indices]
    
    # 打印结果
    print('类别{}的前{}个单词：{}'.format(category, n, top_n_words))