Python中如何基于条件为Pandas DataFrame多列生成指定值

原有代码出错的核心原因是循环内的赋值操作均为对Alloc1、Alloc2整列做覆盖，循环执行到最后一个元素（值为40）时，整列都会被最后一次else分支的结果覆盖，因此出现只有最后一个值生效的问题。pandas处理这类按条件赋值场景时，不建议使用for循环遍历行，向量化操作的效率更高、逻辑也更清晰。

方法1：用loc按条件赋值（逻辑最直观） #

import pandas as pd
import numpy as np

def alloc_update(df):
    # 筛选Number等于99的行
    mask = df['Number'] == 99
    # 给匹配条件的行赋值
    df.loc[mask, 'Alloc1'] = 31
    df.loc[mask, 'Alloc2'] = 3
    # 给不匹配条件的行赋值
    df.loc[~mask, 'Alloc1'] = 0
    df.loc[~mask, 'Alloc2'] = df.loc[~mask, 'Number'] / 2
    return df

调用测试代码：

# 构造原始数据
data = {'ID': ['001', '002', '003', '004'], 'Number': [99, 99, 20, 40], 'Alloc1': [np.NaN, np.NaN, np.NaN, np.NaN], 'Alloc2': [np.NaN, np.NaN, np.NaN, np.NaN]}
df = pd.DataFrame(data)
# 调用封装函数得到结果
final_df = alloc_update(df)

方法2：用np.where实现更简洁的写法 #

def alloc_update(df):
    df['Alloc1'] = np.where(df['Number'] == 99, 31, 0)
    df['Alloc2'] = np.where(df['Number'] == 99, 3, df['Number'] / 2)
    return df

两种方法运行后均可得到预期输出：

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