如何高效为Pandas DataFrame生成判断两列值范围匹配的新列？

问题背景 #

有一个包含A、B两列的大型Pandas DataFrame（1000万行），列值均为0-8的随机整数。需要生成新列Eq：当A的值处于B的值加减指定范围的区间内时，Eq为1，否则为0。

尝试过用apply（包括swifter加速）处理，但速度极慢；直接判断两列相等的方法速度很快，但扩展到范围判断时出现'Series' object cannot be interpreted as an integer错误，现寻求高效实现方案。

测试用DataFrame生成代码：

import pandas as pd
import random
df = pd.DataFrame({'A':[random.randrange(0, 9, 1) for i in range(10000000)],
                   'B':[random.randrange(0, 9, 1) for i in range(10000000)]})

当范围限制为2时，期望结果示例：

A   B   Eq
0   8   3    0
1   3   0    0
2   8   4    0
3   6   5    1
4   8   2    0
...

高效解决方案 #

核心思路是用矢量化操作替代逐行循环（apply本质是逐行处理，效率极低），以下两种方法都能达到和“判断两列相等”相近的速度：

方法1：Pandas原生矢量化布尔判断

直接对Series进行布尔运算，再转换为整数类型。之前报错的原因是误用了range()（该函数不接受Series作为参数），而直接用比较运算符完全兼容Series。

# 设置指定范围
range_limit = 2

# 生成Eq列
df['Eq'] = ((df['A'] >= df['B'] - range_limit) & (df['A'] <= df['B'] + range_limit)).astype(int)

方法2：结合Numpy矢量化操作

对于超大规模数据集，Numpy的底层优化可能带来更快的处理速度：

import numpy as np

range_limit = 2
a_arr = df['A'].to_numpy()
b_arr = df['B'].to_numpy()

df['Eq'] = np.where((a_arr >= b_arr - range_limit) & (a_arr <= b_arr + range_limit), 1, 0)

说明 #

两种方法都是基于矢量化计算，避免了Python层面的逐行循环，能在几秒内完成1000万行数据的处理，效率远高于apply类方法。

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