Python中关联Linkage DataFrame与其他DataFrame的高效方法

问题背景 #

现有两个DataFrame：

• linkage：num1字段包含空格分隔的多个编号，num2为对应标签
• df：num1为单个编号，num2初始值为none

需求：将linkage中匹配的num2值填充到df，匹配条件是df的num1存在于linkage的num1字符串中。

原始代码采用双重列表推导式实现，虽能得到预期结果，但数据量增大时性能急剧下降（测试耗时：CPU times: total: 516 ms，Wall time: 519 ms）：

df.num2 = [linkage.num2[i] for y in df.num1 for i, x in enumerate(linkage.num1) if y in x]

优化方案 #

方案1：拆分映射表后合并（向量化操作） #

先将linkage中num1的多编号拆分为单行记录，生成一对一的映射关系，再通过merge完成填充，利用pandas的向量化操作提升效率：

import pandas as pd

# 原始数据初始化
num1 = ["1111 2222", "3333", "4444 5555 6666", "7777 8888", "9999"]
num2 = ["A1", "A2", "A3", "A4", "A5"] 
linkage = pd.DataFrame({"num1":num1, "num2":num2})
num1 = ["2222", "3333", "5555", "8888", "9999"]
num2 = ['none', 'none', 'none', 'none', 'none']
df = pd.DataFrame({"num1":num1, "num2":num2})

# 拆分linkage的num1字段为多行
linkage_expanded = linkage.assign(num1=linkage['num1'].str.split()).explode('num1')
# 合并并更新num2
df = df.merge(linkage_expanded, on='num1', how='left', suffixes=('_original', ''))
df['num2'] = df['num2'].fillna(df['num2_original'])
df = df.drop('num2_original', axis=1)

print(df)

方案2：构建映射字典批量填充（哈希查找） #

提前构建单个编号到标签的映射字典，通过map方法批量填充，字典的哈希查找时间复杂度为O(1)，在大规模数据下性能最优：

import pandas as pd

# 原始数据初始化（同上）
num1 = ["1111 2222", "3333", "4444 5555 6666", "7777 8888", "9999"]
num2 = ["A1", "A2", "A3", "A4", "A5"] 
linkage = pd.DataFrame({"num1":num1, "num2":num2})
num1 = ["2222", "3333", "5555", "8888", "9999"]
num2 = ['none', 'none', 'none', 'none', 'none']
df = pd.DataFrame({"num1":num1, "num2":num2})

# 构建{单个编号: 对应标签}的映射字典
num_map = {}
for _, row in linkage.iterrows():
    for num in row['num1'].split():
        num_map[num] = row['num2']

# 批量填充df的num2字段
df['num2'] = df['num1'].map(num_map).fillna(df['num2'])

print(df)

性能对比 #

两种方案均避免了原始代码的O(n*m)双重循环：

• 方案1的向量化操作比循环快数倍，代码简洁易读
• 方案2的字典映射在数据量较大时性能更优，耗时可降至毫秒级

内容的提问来源于stack exchange，提问作者Michael S.