如何在指定索引下实现数组相乘并规避中间数组内存占用？

我有两个形状不同的数组，例如：

import numpy as np
matrix = np.arange(5*6*7*8).reshape(5, 6, 7, 8)
vector = np.arange(1, 20, 2)

我想要将matrix的每个元素与vector中的某个元素相乘，然后对最后两个轴求和。为此，我有一个与matrix形状相同的数组Idx，用于指定要使用的vector索引，例如：

Idx = (matrix+np.random.randint(0, vector.size, size=matrix.shape))%vector.size

我知道一种解决方案是：

matVec = vector[Idx]
res = np.sum(matrix*matVec, axis=(2, 3))

或者：

res = np.einsum('ijkl, ijkl -> ij', matrix, matVec)

但我的问题在于数组规模很大，构建matVec既耗时又占用内存。有没有办法绕过这一步仍得到相同结果？

核心思路是避免构建与matrix同形状的matVec，转而先对matrix按Idx的分组进行求和，再与vector对应元素相乘，这样能大幅降低内存占用并提升计算效率。

方法1：循环结合np.bincount（简单易理解） #

对于示例中的4维数组，我们可以将前两个轴作为独立批次，对每个批次内的元素按Idx分组求和，再与vector点乘：

# 重塑数组：前两个轴保留，后两个轴合并为一维
matrix_reshaped = matrix.reshape(matrix.shape[0], matrix.shape[1], -1)
idx_reshaped = Idx.reshape(matrix.shape[0], matrix.shape[1], -1)

res = np.zeros((matrix.shape[0], matrix.shape[1]), dtype=matrix.dtype)
for i in range(matrix.shape[0]):
    for j in range(matrix.shape[1]):
        # 计算当前批次中每个索引对应的matrix元素加权和
        counts = np.bincount(idx_reshaped[i,j], weights=matrix_reshaped[i,j], minlength=vector.size)
        # 加权和与vector点乘得到当前批次结果
        res[i,j] = np.dot(counts, vector)

方法2：向量化实现（高效无循环） #

如果需要更高的效率，可以用np.add.at实现向量化的分组求和，再结合np.einsum完成最终计算：

def vectorized_bincount(idx, weights, minlength):
    # 将数组重塑为（批次数量，元素数量）的二维结构
    idx_flat = idx.reshape(-1, idx.shape[-1])
    weights_flat = weights.reshape(-1, weights.shape[-1])
    # 初始化结果数组：每个批次对应vector长度的求和结果
    out = np.zeros((idx_flat.shape[0], minlength), dtype=weights.dtype)
    # 按索引分组累加权重
    np.add.at(out, (np.arange(idx_flat.shape[0])[:, None], idx_flat), weights_flat)
    # 恢复原批次的形状
    return out.reshape(idx.shape[:-1] + (minlength,))

# 重塑数组
matrix_reshaped = matrix.reshape(matrix.shape[0], matrix.shape[1], -1)
idx_reshaped = Idx.reshape(matrix.shape[0], matrix.shape[1], -1)

# 计算每个索引对应的matrix元素加权和
counts = vectorized_bincount(idx_reshaped, matrix_reshaped, vector.size)
# 加权和与vector相乘并求和，得到最终结果
res = np.einsum('ijk,k->ij', counts, vector)

针对你提供的业务场景，直接应用上述向量化方法替换原计算逻辑即可，代码如下：

import numpy as np

order = 20
dim = 23

listOrder = np.arange(-order, order+1, 1)
N, P = np.meshgrid(listOrder, listOrder)
K = np.arange(-2*dim+1, 2*dim+1, 1)
X = np.arange(-2*dim, 2*dim, 1)

tN = np.einsum('..., p, x -> ...px', N, np.ones(K.shape, dtype=int), np.ones(X.shape, dtype=int))
tP = np.einsum('..., p, x -> ...px', P, np.ones(K.shape, dtype=int), np.ones(X.shape, dtype=int))
tK = np.einsum('..., p, x -> ...px', np.ones(P.shape, dtype=int), K, np.ones(X.shape, dtype=int))
tX = np.einsum('..., p, x -> ...px', np.ones(P.shape, dtype=int), np.ones(K.shape, dtype=int), X)
tL = tK + tX

mini, maxi = -4*dim+1, 4*dim-1
NmPp2L = np.arange(2*mini-2*order, 2*maxi+2*order+1)
Idx = (2*tL+tN-tP) - NmPp2L[0]

np.random.seed(0)
matrix = (np.random.rand(Idx.size) + 1j*np.random.rand(Idx.size)).reshape(Idx.shape)
vector = np.random.rand(np.max(Idx)+1) + 1j*np.random.rand(np.max(Idx)+1)

# 定义向量化分组求和函数
def vectorized_bincount(idx, weights, minlength):
    idx_flat = idx.reshape(-1, idx.shape[-1])
    weights_flat = weights.reshape(-1, weights.shape[-1])
    out = np.zeros((idx_flat.shape[0], minlength), dtype=weights.dtype)
    np.add.at(out, (np.arange(idx_flat.shape[0])[:, None], idx_flat), weights_flat)
    return out.reshape(idx.shape[:-1] + (minlength,))

# 重塑数组，合并最后两个轴
matrix_reshaped = matrix.reshape(matrix.shape[0], matrix.shape[1], -1)
idx_reshaped = Idx.reshape(matrix.shape[0], matrix.shape[1], -1)

# 优化后的计算流程
counts = vectorized_bincount(idx_reshaped, matrix_reshaped, vector.size)
res_optimized = np.einsum('ijk,k->ij', counts, vector)

# 验证与原结果一致性
res_original = np.sum(matrix*vector[Idx], axis=(2, 3))
print(np.allclose(res_optimized, res_original))  # 输出True

内容的提问来源于stack exchange，提问作者Görg