Python中如何避免使用for循环实现Numpy多维数组切片加权求和操作

下面所有方法的运算结果都和你提供的for循环版本完全一致，可以用np.allclose(temp, temp_loop)验证正确性。

方法1：广播+按维度求和（最易理解） #

利用numpy的广播机制，先把一维的y扩展为和X维度匹配的(3,1,1)形状，对应位置相乘后沿第一个维度求和即可：

temp = (y[:, np.newaxis, np.newaxis] * X).sum(axis=0)

方法2：爱因斯坦求和（效率最高） #

用np.einsum直接指定维度映射规则，一步完成乘加操作，不需要额外生成临时数组，运算效率最高：

temp = np.einsum('k, kij -> ij', y, X)

其中参数'k, kij -> ij'的含义是：取y的k维度、X的k/i/j三个维度，对k维度执行乘加运算后输出i/j维度的结果。

方法3：矩阵点乘实现 #

把X的后两个维度拉平为二维数组后和y做矩阵点乘，再reshape回(8,8)的形状即可：

temp = y.dot(X.reshape(X.shape[0], -1)).reshape(X.shape[1], X.shape[2])

结果验证示例 #

你可以用下面的代码验证所有方法和for循环结果一致：

import numpy as np
X = np.random.rand(3,8,8)
y = np.random.rand(3)

# 原始for循环版本
temp_loop = 0
for k in range(3):
    temp_loop = temp_loop + y[k] * X[k,:,:]

# 广播版本
temp_broadcast = (y[:, np.newaxis, np.newaxis] * X).sum(axis=0)
print(np.allclose(temp_loop, temp_broadcast)) # 输出True

# einsum版本
temp_einsum = np.einsum('k, kij -> ij', y, X)
print(np.allclose(temp_loop, temp_einsum)) # 输出True

# 点乘版本
temp_dot = y.dot(X.reshape(3, -1)).reshape(8,8)
print(np.allclose(temp_loop, temp_dot)) # 输出True

内容的提问来源于stack exchange，提问作者Nash J.