基于多键快速在DataFrame间赋值的高效实现方案问询

我有两个DataFrame：

• population：包含两列无序位置数据position1、position2
• keyFrame：包含两列有序键数据key1、key2，以及关联的属性列attr

原实现逻辑：

1. 在population中创建空列assignment
2. 遍历keyFrame的每一行（不修改该DataFrame）
3. 当population的position1与key1、position2与key2匹配，或position1与key2、position2与key1匹配时，将对应行的attr值赋值给population的assignment列

该逻辑可行，但在实际场景（population超50万行，keyFrame超1500条数据）下运行速度极慢。

核心问题：能否一次性实现keyFrame的attr值基于可互换匹配的键对批量赋值给population？

测试代码 #

# Sample code for you to test! Thank you!    
import pandas as pd
import numpy as np

population = pd.DataFrame(data={'position1': [1, 6, 1, 1, 1, 7, 1, 8, 16],
                                'position2': [5, 1, 15, 9, 17, 1, 2, 1, 1]})
keyFrame = pd.DataFrame(data={'key1': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                              'key2': [2, 3, 4, 5, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 17],
                              'attr': [0.79, 0.65, 0.99, 0.03, 0.58, 0.19, 0.53,
                                       0.76, 0.49, 0.46, 0.25, 0.11, 0.22, 0.38, 0.94]})

population['assignment'] = np.NaN  # Step 1
for index, row in keyFrame.iterrows():  # Step 2
    # Step 3
    population['assignment'].loc[((population['position1'] == row['key1']) & (
                population['position2'] == row['key2'])) | (
                (population['position1'] == row['key2']) & (
                population['position2'] == row['key1']))] = row['attr']

最终反馈 #

感谢所有优质建议！这些方案均可行且远快于原实现！

速度测试结果：

1. BeRT2me的方法：18.45秒
2. Jānis Š的方法：21.34秒
3. ouroboros1的方法：26.96秒

注意：这些方案对索引值敏感，请确保重置population和keyFrame的所有索引。

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