Pandas自定义带参数的ifelse函数并应用至多列生成新计算列

问题原因 #

你原来的写法报错是因为apply调用自定义函数传参的语法错误，不能直接写ifelsefunction(row, 'feature1')，需要通过args参数传递额外入参，或者用lambda表达式包装。

方案1：适配最初的简化场景 #

修改后的代码如下：

def ifelsefunction(row, feature):
    if row[feature] >=3:
        return 1
    elif row[feature] ==2:
        return 2
    else:
        return 0

# 单次调用写法：通过args传参
t1['ft1_score'] = t1.apply(ifelsefunction, axis=1, args=('feature1',))
t1['ft2_score'] = t1.apply(ifelsefunction, axis=1, args=('feature2',))
t1['ft3_score'] = t1.apply(ifelsefunction, axis=1, args=('feature3',))

# 批量遍历写法：不用重复写代码
feature_list = ['feature1', 'feature2', 'feature3']
for idx, feat in enumerate(feature_list, 1):
    t1[f'ft{idx}_score'] = t1.apply(ifelsefunction, axis=1, args=(feat,))

方案2：适配补充的实际业务逻辑场景 #

注意np.select的逻辑判断要使用位运算符&替代and，否则会触发数组布尔值判断报错。推荐直接用向量化操作实现，比行遍历apply的执行效率高很多，代码如下：

import numpy as np

# 可根据实际需求替换var2的取值，示例中var2取固定值6
VAR2 = 6

def calc_score(var1, var2):
    mask1 = (var1 >=3) & (var1 < var2)
    mask2 = var1 == 2
    return np.select([mask1, mask2], [var1*0.7, var1*var2], default=0)

# 批量遍历生成新列
feature_list = ['feature1', 'feature2', 'feature3']
for idx, feat in enumerate(feature_list, 1):
    t1[f'ft{idx}_score'] = calc_score(t1[feat], VAR2)

如果实际场景中var2是DataFrame的某一列，直接传入对应列即可：

# 示例：var2取feature3列的值
feature_list = ['feature1', 'feature2']
for idx, feat in enumerate(feature_list, 1):
    t1[f'ft{idx}_score'] = calc_score(t1[feat], t1['feature3'])

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