Python 如何在numpy中对二维阵列上的三重循环进行矢量化？

我是否可以消除此计算中的所有Python循环：

result[i,j,k] = (x[i] * y[j] * z[k]).sum()

其中

x[i]

，

y[j]

，

z[k]

是长度

N

和

x

，

y

，

z

具有长度

A

，

B

，

C

s.t的向量。输出是形状

（A，B，C）

，每个元素都是 三重乘积的和（按元素）

我可以将它从3个循环降到1个循环（代码如下），但我一直在尝试 消除最后一个循环

如有必要，我可以制作

A=B=C

（通过少量填充）

---

如果您使用的是numpy>1.6，那么有一个很棒的

np.einsum

函数：

np.einsum('im,jm,km->ijk',x,y,z)

这相当于您的循环版本。我不确定，一旦达到实际问题中阵列的大小，这将如何公平地提高效率（实际上，当我移动到这些大小时，我的机器上会出现SEGFULT）。对于这类问题，我通常更喜欢的另一种解决方案是使用cython重新编写方法。

在您的案例中，使用非常有意义；但是你可以很容易地用手来做。诀窍是使数组可以相互广播。这意味着重新调整它们的形状，使每个阵列沿其自身的轴独立变化。然后将它们相乘，让

numpy

负责广播；然后沿最后（最右边）轴求和

您可以通过使用切片表示法、

newaxis

值（它等于

None

，因此下面的内容也适用于

None

），以及

sum

接受负轴值这一事实，使这一点更加通用化（使用

-1

表示最后一个，

-2

表示倒数第二个，依此类推）。这样，您就不必知道阵列的原始形状；只要它们的最后一个轴兼容，这将一起广播前三个：

>>> (x[:, numpy.newaxis, numpy.newaxis, :] *
...  y[numpy.newaxis, :, numpy.newaxis, :] *
...  z[numpy.newaxis, numpy.newaxis, :, :]).sum(axis=-1)
array([[[  36,   92,  148,  204],
        [  92,  244,  396,  548],
        [ 148,  396,  644,  892]],

       [[  92,  244,  396,  548],
        [ 244,  748, 1252, 1756],
        [ 396, 1252, 2108, 2964]]])

---

可悲的是，这不是家庭作业问题。事实上，如果有关于“我如何矢量化”这一主题的课程/教科书，我会非常激动！

>>> x = numpy.arange(2 * 4).reshape(2, 4)
>>> y = numpy.arange(3 * 4).reshape(3, 4)
>>> z = numpy.arange(4 * 4).reshape(4, 4)
>>> (x.reshape(2, 1, 1, 4) * 
...  y.reshape(1, 3, 1, 4) *
...  z.reshape(1, 1, 4, 4)).sum(axis=3)
array([[[  36,   92,  148,  204],
        [  92,  244,  396,  548],
        [ 148,  396,  644,  892]],

       [[  92,  244,  396,  548],
        [ 244,  748, 1252, 1756],
        [ 396, 1252, 2108, 2964]]])

>>> (x[:, numpy.newaxis, numpy.newaxis, :] *
...  y[numpy.newaxis, :, numpy.newaxis, :] *
...  z[numpy.newaxis, numpy.newaxis, :, :]).sum(axis=-1)
array([[[  36,   92,  148,  204],
        [  92,  244,  396,  548],
        [ 148,  396,  644,  892]],

       [[  92,  244,  396,  548],
        [ 244,  748, 1252, 1756],
        [ 396, 1252, 2108, 2964]]])