如何向量化实现NumPy数组的逐列加权求和运算？

数组定义 #

import numpy as np

a = np.array([[20, 12,  6],
              [12, 24, 18],
              [ 0, 14, 30]])
b = np.array([1, 0.5])
c = np.array([b ** i for i in range(0, 3)][::-1])
# c 的输出结果：
# array([[1.  , 0.25],
#        [1.  , 0.5 ],
#        [1.  , 1.  ]])

需求说明 #

当使用c的第二列时，需执行如下逐元素乘法后按列求和：

[[20*0.25, 12*0.25, 6*0.25],
 [12*0.5,  24*0.5,  18*0.5],
 [ 0*1,    14*1,    30*1]
].sum(axis=0)

需要对c的每一列都完成上述操作，已知c为一维数组时可通过(a * c[:,None]).sum(axis=0)实现，要求不使用循环，通过向量化方式完成整体运算。

解决方案 #

方法1：矩阵乘法（简洁高效） #

直接对a的转置与c进行矩阵乘法，可一次性得到所有列的运算结果：

result = a.T @ c

执行结果：

print(result)
# 输出：
# [[32.  11. ]
#  [50.  29. ]
#  [54.  40.5]]

方法2：广播机制求和 #

通过扩展维度实现广播，完成逐元素乘法后沿行维度（axis=0）求和：

result = (a[:, :, None] * c[:, None, :]).sum(axis=0)

原理说明：

• a[:, :, None] 将a从(3,3)维度扩展为(3,3,1)
• c[:, None, :] 将c从(3,2)维度扩展为(3,1,2)
• 两者相乘后得到(3,3,2)的数组，沿axis=0求和后得到(3,2)的最终结果，与矩阵乘法结果完全一致。

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