pandas如何基于条件在loc中动态指定列完成高效赋值

完全基于pandas向量化运算实现，无Python层面的逐行循环，3000行数据执行耗时在毫秒级，效率远高于遍历方案：

实现代码 #

import pandas as pd
import numpy as np

# 1. 提取所有月份列的列名
month_cols = [col for col in df.columns if col.startswith('Month ')]

# 2. 单行向量化赋值
df[month_cols] = (
    pd.get_dummies('Month ' + df['col_b'])  # 生成对应月份的哑变量矩阵
    .reindex(columns=month_cols, fill_value=0)  # 对齐原有月份列的顺序和范围
    .mul(df['col_a'], axis=0)  # 哑变量乘以对应行的col_a值，目标位置得到col_a值
    .replace(0, np.nan)  # 非目标位置替换为空值，符合输出要求
)

代码逻辑说明 #

• 因为col_b为字符串类型，直接拼接Month 前缀即可得到对应要赋值的列名
• pd.get_dummies生成的哑变量矩阵天然匹配了「对应列位置为1，其余为0」的掩码效果，和col_a相乘后即可直接得到目标赋值矩阵
• reindex操作保证即使col_b出现不存在的月份值、或者原有月份列存在无匹配的列时，输出结构和原DataFrame完全一致

可选低内存写法 #

如果月份列数量很多，不想生成中间哑变量矩阵占用额外内存，可以用如下写法，仅遍历月份列（最多遍历12次，和数据行数无关）：

import pandas as pd
import numpy as np

month_cols = [col for col in df.columns if col.startswith('Month ')]
# 先初始化所有月份列为空值
df[month_cols] = np.nan
# 仅遍历月份列，列数固定最多12次，不受数据行数影响
for col in month_cols:
    target_num = col.split(' ')[-1]
    df.loc[df['col_b'] == target_num, col] = df.loc[df['col_b'] == target_num, 'col_a']

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