Pandas中基于多列条件的if else逻辑实现问题

问题背景 #

现有如下Pandas DataFrame：

import pandas as pd

df = pd.DataFrame({
    "code": ['sp', 'wh', 'sp'],
    "qty": [20, 30, 10]})

需要创建新列out：当code列值为'sp'时，out取值对应qty列的值；否则取值为0。

用户尝试的代码存在语法错误：

df['out'] = df.apply(lambda x: x['qty']) if x['code'] == 'sp' else 0)

预期输出：

code  qty  out
0    sp   20   20
1    wh   30    0
2    sp   10   10

修正方案 #

方法1：修复apply语法（不推荐，逐行处理效率低） #

你的代码错误在于括号位置错误，if-else逻辑需要放在lambda函数内部，同时必须指定axis=1表示逐行处理：

df['out'] = df.apply(lambda x: x['qty'] if x['code'] == 'sp' else 0, axis=1)

方法2：矢量化操作（推荐，效率远高于逐行处理） #

Pandas优先推荐矢量化操作，避免逐行循环/apply，以下两种方式都可以实现需求：

方式A：使用numpy.where

import numpy as np

df['out'] = np.where(df['code'] == 'sp', df['qty'], 0)

方式B：布尔索引赋值

先将out列初始化为0，再给满足条件的行赋值：

df['out'] = 0
df.loc[df['code'] == 'sp', 'out'] = df['qty']

结果验证 #

执行任意一种方法后，打印df即可得到预期结果：

print(df)
# 输出：
#   code  qty  out
# 0    sp   20   20
# 1    wh   30    0
# 2    sp   10   10

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