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的总和来解决这个问题，但没有成功。

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

对于大型阵列，我们还可以利用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)

坐标的总和

1 个答案: