python - 将 numpy 差异广播到三维

Question

设置

import numpy as np

A = np.array([[1, 2], [2, 3]])
B = np.array([[1, 1], [2, 2], [4, 3]])

A

array([[1, 2],
       [2, 3]])

B

array([[1, 1],
       [2, 2],
       [4, 3]])

我需要取 A 的第一行与 B 的每一行的差异。如果我这样做:

A - B[0]

array([[0, 1],
       [1, 2]])

我只需要为 B 的每一行添加这个。

非矢量化方法是:

np.array([A - B[i] for i in range(B.shape[0])])

array([[[ 0,  1],
        [ 1,  2]],

       [[-1,  0],
        [ 0,  1]],

       [[-3, -1],
        [-2,  0]]])

问题

获得相同 3 维数组的矢量化方法是什么？如果这样更容易，我可以使用 pandas。

Answer 1

最简单的方法是为 B 数组添加一个维度，以便 numpy 正确广播它:

In [15]: A - B[:, np.newaxis]
Out[15]:
array([[[ 0,  1],
        [ 1,  2]],

       [[-1,  0],
        [ 0,  1]],

       [[-3, -1],
        [-2,  0]]])