Python pandas DataFrame无循环实现if else条件赋值的方法

实现思路 #

通过向量化条件判断实现匹配赋值，全程无Python层显式循环，计算效率高，适配不同规模的数据集。

完整代码 #

import pandas as pd
import numpy as np

# 构造原始DataFrame
data = {
  "a": [0, 1, 0,1],
  "b": [0,0,1,1]
}
df = pd.DataFrame(data)

# 构造条件列表和对应取值
conditions = [
    df['a'] == df['b'],
    (df['a'] == 0) & (df['b'] == 1),
    (df['a'] == 1) & (df['b'] == 0)
]
a_values = [1, 0.25, 0.75]
b_values = [1, 0.75, 0.25]

# 生成结果DataFrame
df2 = pd.DataFrame({
    'a': np.select(conditions, a_values),
    'b': np.select(conditions, b_values)
})

结果验证 #

打印df2可得到符合要求的输出：

如果不想额外依赖numpy，也可以用pandas原生的mask方法实现等价逻辑：

df2 = df.copy()
# 相同值场景赋值1
df2[df['a'] == df['b']] = 1
# a=0、b=1场景赋值
mask1 = (df['a'] == 0) & (df['b'] == 1)
df2.loc[mask1, 'a'] = 0.25
df2.loc[mask1, 'b'] = 0.75
# a=1、b=0场景赋值
mask2 = (df['a'] == 1) & (df['b'] == 0)
df2.loc[mask2, 'a'] = 0.75
df2.loc[mask2, 'b'] = 0.25

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