有效地减去不同形状的numpy数组

时间:2015-11-12 17:05:59

标签: python arrays numpy vectorization numpy-broadcasting

使用numpy的优秀广播规则,您可以使用

从形状(5,3)数组v中减去形状(3,)数组X
X - v

结果是一个形状(5,3)数组,其中每行i是差异X[i] - v

有没有办法从w中减去形状(n,3)数组X,以便从整个数组w中减去X的每一行而不明确使用循环?

1 个答案:

答案 0 :(得分:10)

您需要使用None/np.newaxis扩展X的尺寸以形成3D数组,然后按w进行减法。这会为此3D操作带来broadcasting,并产生形状为(5,n,3)的输出。实现看起来像这样 -

X[:,None] - w  # or X[:,np.newaxis] - w

相反,如果所需的排序为(n,5,3),那么您需要扩展w的维度,如此 -

X - w[:,None] # or X - w[:,np.newaxis] 

示例运行 -

In [39]: X
Out[39]: 
array([[5, 5, 4],
       [8, 1, 8],
       [0, 1, 5],
       [0, 3, 1],
       [6, 2, 5]])

In [40]: w
Out[40]: 
array([[8, 5, 1],
       [7, 8, 6]])

In [41]: (X[:,None] - w).shape
Out[41]: (5, 2, 3)

In [42]: (X - w[:,None]).shape
Out[42]: (2, 5, 3)