如何高效按DataFrame类别构建N-grams？优化冗余代码

核心思路 #

借助pandas的groupby分组能力，结合N-gram生成工具（nltk或sklearn）完成批量处理，彻底摆脱逐类别手动循环的冗余写法，同时保证代码的可维护性和执行效率。

方法一：nltk + pandas groupby + Counter（灵活自定义规则） #

适合需要精细化预处理（如过滤停用词、标点、短词）的场景：

import pandas as pd
from nltk.util import ngrams
from nltk.tokenize import word_tokenize
from collections import Counter
import string

# 示例数据集
df = pd.DataFrame({
    'category': ['科技', '科技', '娱乐', '娱乐', '体育'],
    'text': ['人工智能改变生活', '机器学习助力科技发展', '电影票房创新高', '明星红毯造型', '足球联赛夺冠']
})

# 文本预处理函数：分词+过滤无效内容
def preprocess(text):
    tokens = word_tokenize(text)
    # 过滤标点、长度小于2的词，统一转小写
    tokens = [token.lower() for token in tokens if token not in string.punctuation and len(token) >= 2]
    return tokens

# 提取指定类别的Top N N-grams
def get_top_ngrams(tokens_list, n=2, top_k=3):
    all_ngrams = []
    for tokens in tokens_list:
        all_ngrams.extend(ngrams(tokens, n))
    return Counter(all_ngrams).most_common(top_k)

# 分组执行处理
result = df.groupby('category')['text'].apply(
    lambda x: get_top_ngrams(x.apply(preprocess).tolist())
).reset_index(name='top_bigrams')

print(result)

方法二：sklearn CountVectorizer + pandas groupby（高效批量处理） #

面对大规模数据集时，sklearn的向量化实现性能更优，适合快速生成N-gram词频统计：

import pandas as pd
from sklearn.feature_extraction.text import CountVectorizer

# 初始化CountVectorizer，指定生成bigram（可修改ngram_range调整N值）
# 中文场景可传入自定义停用词列表：stop_words=自定义停用词列表
vectorizer = CountVectorizer(ngram_range=(2,2), stop_words=None)

# 提取单类别的Top N N-grams
def get_top_ngrams_sklearn(text_series, top_k=3):
    # 合并当前类所有文本
    combined_text = ' '.join(text_series.tolist())
    # 生成词频矩阵
    X = vectorizer.fit_transform([combined_text])
    # 匹配词汇与对应频次
    vocab = vectorizer.get_feature_names_out()
    counts = X.toarray()[0]
    # 排序取Top K
    top_indices = counts.argsort()[-top_k:][::-1]
    return [(vocab[i], counts[i]) for i in top_indices]

# 分组执行
result_sklearn = df.groupby('category')['text'].apply(get_top_ngrams_sklearn).reset_index(name='top_bigrams')

print(result_sklearn)

关键优化说明 #

• 完全摒弃逐类别循环，利用groupby自动完成分组逻辑
• 两种方法都支持灵活调整N值（修改ngrams的n参数或ngram_range）
• 预处理逻辑统一封装，便于后续修改和复用
• 大规模数据场景优先选择sklearn方案，底层C实现的向量化操作远快于纯Python循环

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