外坐标和

Question

我有两个数组：

# A
[[0 3]
 [2 3]
 [3 1]]

# B
[[2  0]
 [0 -1]
 [0  1]
 [1  0]]

结果应该是坐标方向的外和：

[[2 3],[0 2],[0 4],[1 3]], # [0 3] + each element of B
[[4 3],[2 2],[2 4],[3 3]],
[[5 1],[3 0],[3 2],[4 1]]

我设法用循环解决了它，但是，需要矢量化实现才能处理大型矩阵。我已经用 outer sum 解决了这个问题，但没有成功。

for i in A:
    for j in B:
        print(i+j)

Answer 1

只需将 a 扩展到 3D，保持最后一个轴与 b 的最后一个轴对齐并将它们相加。这将利用 broadcasting 作为矢量化解决方案。要扩展到更高的维度，我们可以使用 np.newaxis/None.

因此，只需执行 -

a[:,None,:] + b[:, :]

跳过用 :'s 指定的冗余最后一个轴，我们将留下 -

a[:,None] + b

示意-

a[:,None,:] :   m x 1 x n
b[:,:]      :       k x n
output      :   m x k x n

对于大型数组，我们还可以利用利用 multi-cores 的 numexpr。为此，我们需要对前面列出的广播方法进行少量修改，例如 -

import numexpr as ne

ne.evaluate('a3D+b',{'a3D':a[:,None]})

计时 -

In [17]: np.random.seed(0)
    ...: a = np.random.rand(1000,3)
    ...: b = np.random.rand(1000,3)

In [18]: %timeit a[:,None] + b
11 ms ± 92.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [19]: %timeit ne.evaluate('a3D+b',{'a3D':a[:,None]})
4.1 ms ± 95.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

外坐标和

Outer sum of coordinates

python

numpy

vectorization