如何在numpy中沿指定轴执行外减操作 实现高维数组配对相减

你所需的功能完全可以通过Numpy原生机制实现，以下是两种常用实现方案：

方法1：维度扩展+广播（最优方案） #

这是性能最高、可读性最好的实现方式，原理是给A新增一个中间维度，触发Numpy的自动广播机制，仅在数组的样本维度做配对，最后一个特征维度直接执行逐元素相减：

import numpy as np

A = np.arange(5 * 2).reshape(-1, 2)
B = np.arange(3 * 2).reshape(-1, 2)

# 写法1：用np.newaxis新增维度
result = A[:, np.newaxis] - B
# 写法2：用expand_dims接口，语义更清晰
# result = np.expand_dims(A, axis=1) - B

# 验证结果完全符合预期
expected_result = np.array([[aa - bb for bb in B] for aa in A ])
assert result.shape == (5, 3, 2)
assert np.all(result == expected_result)

你之前调用np.subtract.outer得到的形状不符合要求，是因为该接口默认会对输入数组的所有维度执行外运算，不会保留最后一个轴作为运算维度，因此会输出(5,2,3,2)的结果。

方法2：基于outer接口调整维度 #

如果你一定要使用ufunc的outer接口实现，也可以对输出做维度转置和重塑，不过性能比广播方案差：

result = np.subtract.outer(A, B).transpose(0, 2, 1, 3).reshape(5, 3, 2)
assert np.all(result == expected_result)

内容的提问来源于stack exchange，提问作者mutableVoid