如何利用NumPy高效实现加权矩阵求和（w₁A₁+w₂A₂+…+wₙAₙ）

Great question! That loop-based approach can get pretty slow when working with large values of n or M—Python loops have overhead that NumPy's optimized C-backed vector operations can avoid entirely. Here are two clean, efficient ways to compute your weighted matrix sum:

Method 1: Broadcasting + Summation #

You can reshape your weight array W to match the dimensions of A (adding singleton dimensions for the matrix axes), then use element-wise multiplication followed by a sum along the first axis. This leverages NumPy's broadcasting rules to handle the scaling automatically:

import numpy as np

# Reshape W to (n, 1, 1) so it broadcasts with A's (n, M, M) shape
weighted_sum = np.sum(W[:, np.newaxis, np.newaxis] * A, axis=0)

The [:, np.newaxis, np.newaxis] (or the shorter [:, None, None]) transforms your (n,1) array into (n,1,1), which aligns perfectly with the first dimension of A. Multiplying them scales each matrix A[i] by W[i], and summing along axis=0 collapses the first dimension to give your final M×M result.

Method 2: Using np.tensordot #

For a more concise approach (and sometimes even better performance), use np.tensordot—it's designed exactly for these kinds of tensor dot products where you want to contract one dimension across two arrays:

weighted_sum = np.tensordot(W, A, axes=([0], [0]))

Here, axes=([0], [0]) tells NumPy to multiply the 0th dimension of W (the n weights) with the 0th dimension of A (the n matrices), then sum over that shared dimension to produce the M×M output.

Example Verification #

Let's test both methods against your loop approach to confirm they give the same result:

n = 3
M = 2
W = np.array([0.5, 1.0, 2.0]).reshape(n, 1)
A = np.array([
    [[1,2],[3,4]],
    [[5,6],[7,8]],
    [[9,10],[11,12]]
])

# Original loop method
temp = np.zeros((M,M), dtype=np.float)
for i in range(n):
    temp += W[i] * A[i]

# Broadcasting method
result_broadcast = np.sum(W[:, None, None] * A, axis=0)

# tensordot method
result_tensordot = np.tensordot(W, A, axes=([0], [0]))

print("Loop Result:\n", temp)
print("Broadcast Result:\n", result_broadcast)
print("Tensordot Result:\n", result_tensordot)

All three outputs will be identical:

Loop Result:
 [[21.5 25. ]
 [28.5 32. ]]
Broadcast Result:
 [[21.5 25. ]
 [28.5 32. ]]
Tensordot Result:
 [[21.5 25. ]
 [28.5 32. ]]

Performance Comparison #

For large datasets, the difference is dramatic. For example, with n=1000 and M=100:

• The loop approach might take several seconds
• Both NumPy methods will finish in milliseconds (since they avoid Python loop overhead entirely)

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