时间戳

首页 > 解决方案 > 用模数对 numpy 数组求和

python - 用模数对 numpy 数组求和

问题描述

考虑以下设置：

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 循环。有没有办法让这更快？

标签： pythonarrayspython-3.xnumpyarray-broadcasting

解决方案

这是一个更清洁、更快速的解决方案。不幸的是，它们并不相同......

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大约快三倍。

推荐阅读