如何将单数据集两列与另一数据集一列映射？求高效方案

以下是两个Pandas数据集的定义及输出：

import pandas as pd

df1 = pd.DataFrame({'id1': 'AAA ABC ACD ADE AEE AFG'.split(),
                   'id2': 'BBB BBC BCD BDE BEE BFG'.split(),})

print(df1)

输出：

id1  id2
0  AAA  BBB
1  ABC  BBC
2  ACD  BCD
3  ADE  BDE
4  AEE  BEE
5  AFG  BFG

df2 = pd.DataFrame({'student_id': 'ABC BBB AAA DEF AEE BEE'.split(),
                    'center': '11 22 33 44 55 66'.split()})

print(df2)

输出：

student_id center
0        ABC     11
1        BBB     22
2        AAA     33
3        DEF     44
4        AEE     55
5        BEE     66

需求 #

将df1的id1、id2列分别与df2的student_id列映射，仅保留id1和id2均存在于student_id中的行，并生成对应的映射列（包含student_id和center信息）。

当前方法的问题 #

目前用多次merge的方式在小数据集能得到预期结果：

map1 = df1.merge(df2, left_on='id1', right_on='student_id').drop(columns=['id2'])
map2 = df1.merge(df2, left_on='id2', right_on='student_id')
map1.merge(map2, on='id1') 

但大数据集下（比如map1有10万行、map2有7万行）效率极低，结果行数甚至接近100万，设置索引也没有改善，需要更优实现方式。

预期输出 #

id1  id2 student_id1 center_1 student_id2 center_2
0  AAA  BBB         AAA       33         BBB       22 # id1和id2都在student_id中，对应df2的取值
1  AEE  BEE         AEE       55         BEE       66 # id1和id2都在student_id中，对应df2的取值

思路 #

先筛选出df1中id1和id2都存在于df2的student_id中的行，再通过字典映射或精准merge的方式添加对应信息，避免不必要的笛卡尔积操作（这是原方法大数据集行数暴增的核心原因）。

方案一：字典映射法（最高效） #

# 提取df2中有效student_id集合，用于快速判断
valid_ids = set(df2['student_id'])

# 筛选df1中符合条件的行
filtered_df1 = df1[(df1['id1'].isin(valid_ids)) & (df1['id2'].isin(valid_ids))]

# 将df2转为字典，实现O(1)时间复杂度的映射
center_map = df2.set_index('student_id')['center'].to_dict()

# 添加映射列
filtered_df1['student_id1'] = filtered_df1['id1']
filtered_df1['center_1'] = filtered_df1['id1'].map(center_map)
filtered_df1['student_id2'] = filtered_df1['id2']
filtered_df1['center_2'] = filtered_df1['id2'].map(center_map)

# 调整列顺序匹配预期输出
result = filtered_df1[['id1', 'id2', 'student_id1', 'center_1', 'student_id2', 'center_2']]
print(result)

方案二：先筛选再merge（更贴合Pandas风格） #

# 提取有效id集合
valid_ids = set(df2['student_id'])
# 筛选df1中符合条件的行
filtered_df1 = df1[(df1['id1'].isin(valid_ids)) & (df1['id2'].isin(valid_ids))]

# 两次精准merge，仅对筛选后的行操作
result = filtered_df1.merge(df2, left_on='id1', right_on='student_id', suffixes=('_1', '')) \
                     .merge(df2, left_on='id2', right_on='student_id', suffixes=('', '_2')) \
                     .rename(columns={'student_id': 'student_id1', 'student_id_2': 'student_id2',
                                      'center': 'center_1', 'center_2': 'center_2'}) \
                     [['id1', 'id2', 'student_id1', 'center_1', 'student_id2', 'center_2']]
print(result)

效率说明 #

• 两种方案都先做了筛选，大幅减少后续操作的数据量
• 字典映射法的时间复杂度为O(n)，适合超大规模数据集
• 先筛选再merge的方式避免了原方法中无效行带来的笛卡尔积，性能提升显著

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