高效实现DataFrame单行与全量DataFrame近似匹配并生成新列

需求说明 #

需要给DataFrame df新增一列B，其中每行的B元素是另一DataFrame ddf的A列中，与df对应行A值接近（定义为处于原数值的0.8~1.2倍区间内）的所有值构成的列表。

输入示例 #

df
Out[19]: 
    A
0  10
1  20
2  30
3  40

ddf
Out[20]: 
    A
0   9
1  10
2  31
3  50
4  60

预期输出 #

df
Out[22]: 
    A        B
0  10  [9, 10]
1  20       []
2  30     [31]
3  40       []

原有代码的问题 #

原代码df['B'] = ddf[ddf['A'].between(df['A']*0.8, df['A']*1.2)]['A']存在两个核心问题：

• 当df与ddf行数不同时，会因维度不匹配直接报错；
• 即使行数相同，也只能实现逐行一对一匹配，无法让df的每行去匹配ddf的全表数据。

解决方案 #

方法一：apply遍历（简洁直观，中小数据量适用） #

直接对df的A列每行值进行遍历，筛选ddf中符合区间条件的值并转为列表：

df['B'] = df['A'].apply(lambda x: ddf.loc[ddf['A'].between(x*0.8, x*1.2), 'A'].tolist())

方法二：向量化实现（高效，大数据量推荐） #

利用numpy广播生成匹配掩码矩阵，避免循环遍历，提升效率：

import numpy as np

# 生成二维掩码：每行对应df的一个A值，每列对应ddf的一个A值，标记是否符合区间条件
mask = (ddf['A'].values >= df['A'].values[:, np.newaxis] * 0.8) & \
       (ddf['A'].values <= df['A'].values[:, np.newaxis] * 1.2)

# 逐行提取符合条件的ddf['A']值，转为列表
df['B'] = [ddf['A'][row_mask].tolist() for row_mask in mask]

两种方法都能实现需求，方法二更适合处理大规模数据集，避免apply的循环开销。

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