如何加速Pandas中自定义的keep_inum函数？

我拥有三个数值列表list1、list2、list3，以及如下结构的Pandas DataFrame：

我编写了如下自定义函数keep_inum：

def keep_inum(row):
    if len(row) != 0:
        if int(row['inum']) in list1:
            if row['DESC_1'] == 1:
                return row['recs']
            else:
                return ''
        elif int(row['inum']) in list2:
            if row['DESC_1'] == 2:
                return row['recs']
            else:
                return ''
        elif int(row['inum']) in list3:
            if row['DESC_1'] == 3:
                return row['recs']
            else:
                return ''
        else:
            return row['recs']
    else:
        pass

并通过df['recs'] = df.apply(keep_inum, axis = 1)将该函数应用到DataFrame上，请问如何加速这个自定义函数？

Pandas的apply逐行处理效率极低，数据量越大差距越明显，下面是几种高效替代方案：

1. 先优化成员判断效率 #

列表的in操作时间复杂度是O(n)，换成集合后是O(1)，先把三个列表转成集合：

set1 = set(list1)
set2 = set(list2)
set3 = set(list3)

2. 用Pandas矢量化操作替代逐行循环 #

这是效率最高的方案，直接用布尔索引批量处理，完全避免逐行遍历：

# 先把inum转为整数类型（如果原数据不是的话）
df['inum'] = df['inum'].astype(int)

# 复制原recs作为初始结果
result = df['recs'].copy()

# 定义需要将recs置空的条件
mask1 = df['inum'].isin(set1) & (df['DESC_1'] != 1)
mask2 = df['inum'].isin(set2) & (df['DESC_1'] != 2)
mask3 = df['inum'].isin(set3) & (df['DESC_1'] != 3)

# 满足任一条件的位置置空
result[mask1 | mask2 | mask3] = ''

# 赋值回原DataFrame
df['recs'] = result

3. 用numpy.where简化逻辑 #

如果喜欢更紧凑的写法，可以用numpy的where函数组合条件，效率和矢量化操作一致：

import numpy as np

df['inum'] = df['inum'].astype(int)
set1, set2, set3 = set(list1), set(list2), set(list3)

df['recs'] = np.where(
    (df['inum'].isin(set1) & (df['DESC_1'] != 1)) |
    (df['inum'].isin(set2) & (df['DESC_1'] != 2)) |
    (df['inum'].isin(set3) & (df['DESC_1'] != 3)),
    '',
    df['recs']
)

4. 映射规则法（可选） #

如果需要更清晰的规则映射，可以先构建inum到预期DESC_1的字典，再做判断：

# 构建inum对应的预期DESC_1值
desc_map = {}
desc_map.update({num: 1 for num in list1})
desc_map.update({num: 2 for num in list2})
desc_map.update({num: 3 for num in list3})

# 生成每行inum对应的预期DESC_1，无匹配则为None
df['expected_desc'] = df['inum'].astype(int).map(desc_map)

# 当预期值存在且与实际DESC_1不符时置空
df['recs'] = np.where(
    df['expected_desc'].notna() & (df['expected_desc'] != df['DESC_1']),
    '',
    df['recs']
)

# 清理临时列
df.drop('expected_desc', axis=1, inplace=True)

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