如何使用Pandas对关联dummy变量为1的数值列按行求和

解决方案 #

核心思路是利用pandas向量化运算，将数值列与对应dummy列逐元素相乘后按行求和，全程不需要显式遍历列，大样本下效率极高。

代码实现 #

import pandas as pd
import numpy as np

# 自行定义数值列列表，列数再多也只需修改此处
val_cols = ['A', 'B', 'C']
# 自动生成对应dummy列名，保证顺序与数值列一一对应
dummy_cols = [f"{col}_dummy" for col in val_cols]

# 核心计算逻辑：对应列相乘后按行求和
df['ABC_sum'] = (df[val_cols] * df[dummy_cols].values).sum(axis=1)

# 若全0行需要返回NaN而非0，使用下方代码替换上面的赋值语句
# df['ABC_sum'] = (df[val_cols] * df[dummy_cols].values).sum(axis=1).replace(0, np.nan)

方案优势 #

• 全程为向量化运算，无Python层面的循环，比逐列遍历判断效率高1~2个数量级，支持上百列、千万行级别的大规模数据集
• 无需硬编码列判断逻辑，仅需维护val_cols数值列列表即可，列数变动时核心代码无需修改
• 内存占用低，运算过程中不会生成额外的大体积中间对象

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