请考虑以下设置：
import numpy as np
import itertools as it
A = np.random.rand(3,3,3,16,3,3,3,16) # sum elements of A to arrive at...
B = np.zeros((4,4)) # a 4x4 array (output)
我想对一个大数组'A'求和，但是以一种非常特定的方式。 'A'的形状为（x，x，x，16，x，x，x，16），其中'x'是一些整数。 理想的结果是一个4x4的矩阵“ B”，我可以通过如下所示的for循环进行计算：
%%timeit
for x1,y1,z1,s1 in it.product(range(3), range(3), range(3), range(16)):
for x2,y2,z2,s2 in it.product(range(3), range(3), range(3), range(16)):
B[s1%4, s2%4] += A[x1,y1,z1,s1,x2,y2,z2,s2]
>> 134 ms ± 1.27 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
其中B的元素是两个轴的“模4”，在“ A”中的那个维度上有16个元素，这里用s1和s2索引。
如何通过广播或其他方式实现相同目标？显然，对于较大的“ x”（尺寸为“ A”），for循环的计算时间将成指数增长，这是不理想的。
编辑：
C = np.zeros((4,4))
for i,j in it.product(range(4), range(4)):
C[i,j] = A[:,:,:,i::4,:,:,:,j::4].sum()
这似乎也起作用。但仍涉及1个for循环。有没有办法使它更快？
答案 0 :(得分：2)
这是一个更干净，更快的解决方案。不幸的是，它们并不相同...
def clean(A):
return A.reshape(4*n*n*n, 4, 4*n*n*n, 4).sum(axis=(0, 2))
def fast(A):
return np.bincount(np.tile(np.arange(16).reshape(4, 4), (4, 4)).ravel(), A.sum((0,1,2,4,5,6)).ravel(), minlength=16).reshape(4, 4)
在n==6 fast处快大约三倍。