有没有办法仅在维度的子集上使用 numpy.outer？

Question

我有一个数组数组，像这样：

import numpy as np
a = np.array([[1, 2, 3], [4, 5, 6]])

我想计算数组元素之间的成对差异，例如：

[[[ 0,  0,  0], [-3, -3, -3]],
 [[ 3,  3,  3], [ 0,  0,  0]]

我的第一个想法是使用np.subtract.outer(a, a)，但这并没有达到我想要的效果——它在数组中的一层太深了。我可以看到我需要的数字在np.subtract.outer(a, a) 的输出中，但是我实际使用的数组非常大，我没有足够的内存来分配结果。

谢谢！

Answer 1

你可以简单地使用广播来解决这个问题。

a[:, None, :] - a[None, :, :]

给你你想要的。