迭代具有两个隐式循环的数组

Question

是否可以隐式迭代具有两个索引的数组？ 这是我想做的一个非常简单的例子：

import numpy as np

x = np.arange(3)
y = np.zeros(3)

for i in range(3):
    y[i] = np.sum(x - x[i])

有一个隐式循环（总和）和一个显式循环（for i in range(3)）...是否有可能有一个完全隐式的版本？

Answer 1

以下应该有效：

y = np.sum(x - x[:,None], axis=1)

此解决方案使用 numpy 的广播工具。首先，我使用x[:,None] 将形状为(N,) 的x 重铸为(N,1)。您可能还会看到它写成x[:,np.newaxis]。

x - x[:,None] 创建一个(N,N) 数组，其元素为tmp_{i,j} = x_i - x_j。然后，我只需使用np.sum 中的参数axis=1 对各行求和。

见：

In [13]: y = np.zeros(10)

In [14]: x = np.random.normal(size=(10,))

In [15]: for i in range(10):
        y[i] = np.sum(x - x[i])
   ....:

In [16]: y
Out[16]:
array([  7.99781458,   4.15114434, -17.24655912, -20.35606168,
        -5.0211756 ,   7.52062868,   8.2501526 ,   3.90397351,
        10.18746451,   0.61261819])

In [17]: np.sum(x - x[:,None], 1)
Out[17]:
array([  7.99781458,   4.15114434, -17.24655912, -20.35606168,
        -5.0211756 ,   7.52062868,   8.2501526 ,   3.90397351,
        10.18746451,   0.61261819])

In [18]: np.allclose(y, np.sum(x - x[:,None], 1))
Out[18]: True

时间安排： 需要指出的是，使用 numpy 提供的工具来操作数组通常比使用标准 Python 构造快得多：

In [48]: x = np.random.normal(size=(100,))

In [49]: %timeit y = np.array([sum(x - k) for k in x])
100 loops, best of 3: 6.86 ms per loop

In [67]: %timeit y = np.array([np.sum(x - k) for k in x])
1000 loops, best of 3: 1.54 ms per loop

In [50]: %timeit np.sum(x - x[:,None], 1)
10000 loops, best of 3: 59 µs per loop

In [51]:

In [51]: x = np.random.normal(size=(1000,))

In [52]: %timeit y = np.array([sum(x - k) for k in x])
1 loops, best of 3: 592 ms per loop

In [72]: %timeit y = np.array([np.sum(x - k) for k in x])
100 loops, best of 3: 17.2 ms per loop

In [53]: %timeit np.sum(x - x[:,None], 1)
100 loops, best of 3: 8.67 ms per loop

Answer 2

如果可能，您应该始终尝试在计算机科学之前使用数学。你的表达式y[i] = np.sum(x - x[i]) 可以用一点代数重写为y[i] = np.sum(x) - x.size * x[i]。这清楚地表明您可以在没有任何循环的情况下重写代码：

y = np.sum(x) - x.size * x

显然，对于大型数组，它的运行速度比@JoshAdel 的解决方案快得多，对于大小为 1000 的输入，它的运行速度要快 x400：

>>> x = np.random.normal(size=(1000,))
>>> np.allclose(np.sum(x - x[:,None], 1), np.sum(x) - x.size * x)
True

%timeit np.sum(x - x[:,None], 1)
100 loops, best of 3: 6.33 ms per loop

%timeit np.sum(x) - x.size * x
100000 loops, best of 3: 16.5 us per loop