如何在NumPy中无需创建三维数组实现指定复杂求和？

要计算 C_ij = sum_k exp(A_ij*B_k) 且避免生成内存密集的三维数组，你可以采用以下两种高效方法：

方法一：循环累加（最优内存友好方案） #

这种方法通过逐元素遍历B，每次仅生成一个与A同形状的二维数组并累加，全程只占用O(N²)内存，彻底规避三维数组的创建：

import numpy as np

# 示例数据
N = 1000
A = np.random.rand(N, N)
B = np.random.rand(N, 1)

# 初始化结果数组
C = np.zeros_like(A)

# 遍历B的每个元素，累加exp(A*b)
for b_val in B.flat:
    C += np.exp(A * b_val)

原理说明 #

每次循环中，np.exp(A * b_val) 生成一个N×N的二维数组，直接累加到结果C中。由于不需要存储所有中间的三维数据，内存占用仅为原方法的1/N，对大数组场景极为友好。

方法二：展平广播求和（内存占用略优于三维数组） #

如果你的内存能容纳N²×N的二维数组，可以通过展平A来将三维运算转化为二维运算，本质上是用二维数组替代三维数组，内存占用量级相同但结构更紧凑：

# 展平A为(N²,1)，与B转置后的(1,N)广播相乘，exp后求和再重塑形状
C = np.exp(A.reshape(-1, 1) * B.T).sum(axis=1).reshape(A.shape)

补充说明 #

你提到的np.einsum()确实无法直接处理包含exp的非线性求和逻辑，因为它仅针对线性的乘积求和操作。因此上述两种方法是更适合的替代方案。

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