Pandas DataFrame如何按索引提取数组元素生成新列实现词干还原

问题说明 #

现有一个包含3列的Pandas DataFrame：

• stem_word：提取出的词干化目标词汇列表
• stem_tokens：整段文本完成词干化处理的全量分词结果
• tokens：同一段文本未做词干处理的原始分词结果

初始数据构造代码：

import pandas as pd
data = [{'stem_word': ['cat', 'run','today'], 'stem_tokens':  ['cat', 'run', 'today', 'in', 'the' ,'morning', 'very', 'quick'], 'tokens': ['cat', 'running', 'today', 'in' ,'the', 'morning', 'very', 'quickly']}]
df = pd.DataFrame(data)

初始数据结构如下：

需要完成两步操作，最终把词干化词汇还原为对应的原始未词干词汇，存入新列new：

1. 定位stem_word中每个词汇在stem_tokens列里的对应索引位置
2. 用拿到的索引列表，从tokens列提取对应位置的元素组成新列

目标最终效果：

实现方法 #

直接逐行处理即可，用字典存储词干到索引的映射，查询效率比遍历列表高很多：

def match_original_word(row):
    # 为当前行的词干分词构建 {词干: 索引} 映射
    stem_index_map = {token: idx for idx, token in enumerate(row['stem_tokens'])}
    # 按stem_word的顺序拿到匹配索引，再从原始分词中取对应值
    match_index = [stem_index_map[w] for w in row['stem_word']]
    return [row['tokens'][i] for i in match_index]

df['new'] = df.apply(match_original_word, axis=1)

运行后得到的结果和目标效果完全一致。

边界情况处理 #

如果数据中存在stem_word的词干不在stem_tokens中的情况，给列表推导加一层存在性判断即可，避免触发KeyError：

match_index = [stem_index_map[w] for w in row['stem_word'] if w in stem_index_map]

这个方法的前提是stem_tokens和tokens的分词位置完全一一对应，这也是词干化处理的默认输出规则，不会出现匹配错位问题，比单独给每个原始词重复作词干匹配的效率高很多。

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