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

Posted 2023-02-23

【中文标题】有效地减去不同形状的numpy数组

【英文标题】：Subtracting numpy arrays of different shape efficiently

【发布时间】：2016-02-14 02:47:21

【问题描述】：

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

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

X - v

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

有没有办法从 X 中减去形状 (n,3) 数组 w，以便在不显式使用循环的情况下从整个数组 X 中减去 w 的每一行？

【问题讨论】：

输出数组的形状是什么？

我们需要从一个形状 (5,3) 数组中减去 n 个形状 (3,) 向量，所以结果将是一个形状 (5,3,n) 数组。

我认为应该是(5,n,3)，不是吗？因为最后一个轴是行。

事实上排序并不重要，因为我可以使用swapaxes 来实现所需的排序（在我的情况下为 (n,5,3)）。

【参考方案1】：

您需要将X 的维度扩展为None/np.newaxis 以形成一个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)

【讨论】：

这就是我需要的！简单地扩展维度的好主意。谢谢！