如何用NumPy向量化实现4维ndarray的自定义求和运算？

给定形状为(N,N,N,N)的4维NumPy数组A，需要计算形状为(N,N)的2维数组B，其中每个元素B[n,m]是对所有0 ≤ i < n、0 ≤ j < m的A[i,j,n-i,m-j]求和——你说得没错，这确实是一种类卷积的单数组求和操作，完全可以用NumPy实现向量化，避免低效的Python嵌套循环。

核心思路 #

原循环的本质是：对每个(n,m)，求和所有满足i + k = n、j + l = m（其中k = n-i ≥1、l = m-j ≥1）的A[i,j,k,l]元素。我们可以通过索引广播+分组求和或者**爱因斯坦求和（einsum）**来实现向量化。

方法1：索引广播+bincount分组求和 #

这种方法通过生成所有合法的索引组合，将4维数组的元素映射到(n,m)的二维索引上，再用bincount高效求和：

import numpy as np

# 示例输入
N = 5
A = np.random.rand(N, N, N, N)

# 生成所有可能的索引
i, j, k, l = np.ogrid[:N, :N, :N, :N]

# 计算对应的(n,m)，并筛选合法条件：n<N, m<N, k≥1, l≥1
n = i + k
m = j + l
valid_mask = (n < N) & (m < N) & (k >= 1) & (l >= 1)

# 提取合法的元素和对应的(n,m)索引
valid_A = A[valid_mask]
valid_n = n[valid_mask]
valid_m = m[valid_mask]

# 将二维索引转为一维，用bincount求和后重塑为二维数组
flat_idx = valid_n * N + valid_m
B_vec = np.bincount(flat_idx, weights=valid_A, minlength=N*N).reshape(N, N)

方法2：用广播简化内层求和逻辑 #

如果不需要彻底去掉外层循环，这种方法保留n,m的外层遍历，但将内层i,j的求和完全转为NumPy向量化操作，效率远高于原Python嵌套循环：

import numpy as np

N = 5
A = np.random.rand(N, N, N, N)

B_vec = np.zeros((N, N))
for n in range(N):
    for m in range(N):
        # 生成合法的i,j范围，用广播构造索引网格
        i_range = np.arange(n)
        j_range = np.arange(m)
        # 直接对对应位置的A元素求和
        B_vec[n, m] = A[i_range[:, None], j_range, n - i_range[:, None], m - j_range].sum()

验证正确性 #

可以对比原循环的结果和向量化结果，确认一致性：

# 原循环实现
B_loop = np.zeros((N, N))
for n in range(N):
    for m in range(N):
        B_loop[n, m] = sum(A[i,j,n-i,m-j] for i in range(n) for j in range(m))

# 验证结果是否一致
print(np.allclose(B_loop, B_vec))  # 输出True表示结果完全一致

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