如何高效计算文本数据中各分类下的Top N高频词

针对你遇到的分组统计高频词时内存不足、分组求和过慢的问题，我分享几个实战性的优化方向，覆盖稀疏矩阵、高效统计、并行计算三种思路：

一、用稀疏矩阵直接分组聚合（推荐，兼顾速度和内存） #

你的初始思路提到了稀疏矩阵，其实可以全程基于scipy稀疏矩阵完成分组求和，避免将稀疏矩阵转为稠密数组（这是内存爆炸的关键），而且矩阵运算的速度远快于Counter逐组累加。这个方案还能保留你初始需求的n-gram支持（原优化方案只处理了单token）：

from sklearn.datasets import fetch_20newsgroups
import pandas as pd
import random
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.preprocessing import LabelEncoder
import scipy.sparse as sp

# 加载示例数据
data = fetch_20newsgroups(subset='train')
data = pd.DataFrame({'text': data.data[:1000]})
CAT1 = [f'cat1_{i}' for i in range(3)]
CAT2 = [f'cat2_{i}' for i in range(10)]
data['cat1'] = [random.choice(CAT1) for _ in range(data.shape[0])]
data['cat2'] = [random.choice(CAT2) for _ in range(data.shape[0])]

# 1. 生成词频稀疏矩阵（支持1-3gram，自动过滤低频词）
vectorizer = CountVectorizer(ngram_range=(1, 3), min_df=5, max_df=1.0)
X = vectorizer.fit_transform(data['text'])  # 稀疏矩阵，内存占用仅为稠密数组的几分之一

# 2. 生成唯一分组ID（cat1+cat2的组合编码）
data['group_key'] = data['cat1'] + '|' + data['cat2']
le = LabelEncoder()
data['group_id'] = le.fit_transform(data['group_key'])
n_groups = data['group_id'].nunique()

# 3. 构建分组指示矩阵（每行仅对应分组的位置为1，稀疏格式）
group_indicator = sp.csr_matrix(
    ([1]*len(data), (range(len(data)), data['group_id']))
)

# 4. 矩阵乘法实现高效分组求和：(n_groups, 1000) × (1000, n_features) = (n_groups, n_features)
grouped_word_counts = group_indicator.T.dot(X)

# 5. 转换为DataFrame并提取TopN词
group_names = le.inverse_transform(range(n_groups))
grouped_df = pd.DataFrame.sparse.from_spmatrix(
    grouped_word_counts,
    index=group_names,
    columns=vectorizer.get_feature_names_out()
)

# 拆分分组键为cat1、cat2
grouped_df = grouped_df.reset_index().rename(columns={'index': 'group_key'})
grouped_df[['cat1', 'cat2']] = grouped_df['group_key'].str.split('|', expand=True)
grouped_df.drop('group_key', axis=1, inplace=True)

# 提取每个分组的Top10词
def extract_top_words(row, top_n=10):
    # 稀疏转稠密仅针对当前行，内存压力小
    word_counts = row.iloc[:, 2:].sparse.to_dense()
    top_words = word_counts.nlargest(top_n).reset_index()
    top_words.columns = ['WORDS', 'Value']
    return top_words

topwords = grouped_df.apply(extract_top_words, axis=1).explode().reset_index(drop=True)
# 拼接cat1、cat2列
topwords = pd.concat(
    [grouped_df[['cat1', 'cat2']].repeat(10).reset_index(drop=True), topwords],
    axis=1
)

优势： #

• 全程稀疏矩阵操作，内存占用极低，避免了X.toarray()的内存爆炸问题
• 分组求和用矩阵乘法实现，时间复杂度远低于Counter逐元素累加，速度提升明显
• 完整保留n-gram统计能力，符合你初始需求

二、优化现有Counter方案：合并文本后一次性统计 #

你当前的Counter方案慢在每行生成Counter再逐组求和，可以改为先按分组合并所有文本，再一次性统计词频，减少Counter的累加次数：

from collections import Counter
import itertools

# 生成连续n-gram的工具函数（替代split的单token）
def generate_ngrams(text, ngram_range=(1, 3)):
    tokens = text.split()
    ngrams_list = []
    for n in range(ngram_range[0], ngram_range[1]+1):
        if len(tokens) >= n:
            ngrams_list.extend([' '.join(tokens[i:i+n]) for i in range(len(tokens)-n+1)])
    return ngrams_list

# 按分组合并文本，一次性统计词频
grouped_data = data.groupby(['cat1', 'cat2'])['text'].apply(lambda x: ' '.join(x))
# 对合并后的文本统计词频
grouped_counts = grouped_data.apply(lambda x: Counter(generate_ngrams(x)))
# 提取Top3词并展开
topwords = grouped_counts.apply(lambda x: x.most_common(3)).explode().reset_index()
# 拆分词和计数
topwords[['WORDS', 'Value']] = pd.DataFrame(topwords[0].tolist(), index=topwords.index)
topwords.drop(0, axis=1, inplace=True)

优势： #

• 代码简洁易读，无需学习稀疏矩阵操作
• 每个分组仅做一次Counter统计，避免了大量Counter对象的累加操作，速度比原优化方案提升数倍

三、Dask并行计算：处理超大规模数据集 #

如果你的数据量已经超过单节点内存，Dask是很好的选择，它能自动拆分数据并并行执行任务，支持超过内存的大数据处理：

import dask.dataframe as dd
from dask.distributed import Client

# 启动本地Dask集群（也可连接分布式集群）
client = Client()

# 将Pandas DataFrame转为Dask DataFrame，自动拆分分区
ddata = dd.from_pandas(data, npartitions=4)

# 定义词频统计函数（和之前一致）
def count_ngrams(text):
    tokens = text.split()
    ngrams_list = []
    for n in range(1, 4):
        if len(tokens) >= n:
            ngrams_list.extend([' '.join(tokens[i:i+n]) for i in range(len(tokens)-n+1)])
    return Counter(ngrams_list)

# 并行计算每行词频
ddata['text_count'] = ddata['text'].apply(count_ngrams, meta=object)
# 并行分组求和
grouped_counts = ddata.groupby(['cat1', 'cat2'])['text_count'].sum().compute()

# 后续提取Top词的逻辑和之前一致
topwords = grouped_counts.apply(lambda x: x.most_common(3)).explode().reset_index()
topwords[['WORDS', 'Value']] = pd.DataFrame(topwords[0].tolist(), index=topwords.index)
topwords.drop(0, axis=1, inplace=True)

优势： #

• 支持TB级别的超大数据集，无需担心内存不足
• 自动并行执行，利用多核CPU提升速度
• API和Pandas高度兼容，学习成本低

方案选择建议 #

• 若需要保留n-gram统计且数据量中等：优先选择稀疏矩阵方案，速度和内存最优
• 若仅需单token统计且数据量较小：选择优化后的Counter方案，代码简洁
• 若数据量超大（超过单节点内存）：选择Dask并行方案

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