有没有办法只在一个维度的子集上使用 numpy.outer?

Is there a way to use numpy.outer on only a subset of dimensions?

我有一个数组数组,像这样:

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) 的输出中,但我实际使用的数组非常大,我没有足够的内存来分配结果。

谢谢!

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

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

给你你想要的。