基于条件选择mul数组与DataFrame执行元素级乘法的实现方案问询

数据定义 #

import numpy as np
import pandas as pd

df1 = pd.DataFrame({'A': [100, 200], 'B': [50, 150], 'C': [40, 120]})
#      A    B    C
# 0  100   50   40
# 1  200  150  120
df2 = pd.DataFrame({'A': [5, 1], 'B': [3, 25], 'C': [10, 11]})
#    A   B   C
# 0  5   3  10
# 1  1  25  11

mul1 = np.array([.1, .2, .3])
mul2 = np.array([1.1, 1.2, 1.3])

需求说明 #

根据df2的列值条件，选择mul1或mul2对应位置的系数与df1逐元素相乘：

• A列：当df2['A'] > 3时用mul1[0]，否则用mul2[0]
• B列：当df2['B'] > 12时用mul1[1]，否则用mul2[1]
• C列：当df2['C'] > 15时用mul1[2]，否则用mul2[2]

期望输出：

A    B   C
0   10   60  52
1  220   30  156

高效实现方法 #

利用numpy和pandas的向量化操作实现，避免循环，保证执行效率：

# 生成各列的条件掩码（布尔值矩阵）
mask = df2 > [3, 12, 15]
# 根据掩码选择对应系数，自动广播为与df1同形状的矩阵
coeffs = np.where(mask, mul1, mul2)
# 逐元素相乘得到结果
result = df1 * coeffs

print(result)

代码解释 #

1. 构造条件掩码：df2 > [3,12,15]借助广播特性，直接对df2每列与对应阈值比较，生成和df1形状一致的布尔值DataFrame。
2. 选择系数矩阵：np.where根据掩码的布尔值，选择mul1或mul2的对应元素，一维的系数数组会自动广播为二维矩阵，匹配df1的形状。
3. 执行乘法：pandas的DataFrame与numpy矩阵直接相乘时，会自动按元素对应计算，全程向量化操作，效率远高于循环遍历。

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