如何对Numpy二维数组单维度执行外求和？及行内元素相减生成3D数组

问题1：仅对二维数组的单个维度执行外求和操作 #

要实现单个维度的外求和，核心是利用广播机制或ufunc.outer配合维度调整，针对行/列维度的需求分别实现：

场景1：对每行（axis=1）执行外求和 #

每行内所有元素两两求和，最终得到(n, m, m)形状的三维数组（n为行数，m为列数）：

import numpy as np

arr = np.array([[1, 2, 3], [4, 5, 6]])
# 扩展维度后广播求和
row_outer_sum = arr[:, None, :] + arr[:, :, None]
print(row_outer_sum)
# 输出：
# [[[ 2  3  4]
#   [ 3  4  5]
#   [ 4  5  6]]
#
#  [[ 8  9 10]
#   [ 9 10 11]
#   [10 11 12]]]

场景2：对每列（axis=0）执行外求和 #

每列内所有元素两两求和，最终得到(n, n, m)形状的三维数组：

col_outer_sum = arr[None, :, :] + arr[:, None, :]
print(col_outer_sum)
# 输出：
# [[[ 2  4  6]
#   [ 5  7  9]]
#
#  [[ 5  7  9]
#   [ 8 10 12]]]

如果偏好ufunc.outer，可针对每行单独处理后堆叠：

row_outer_sum = np.stack([np.add.outer(row, row) for row in arr], axis=0)

问题2：每行内元素两两相减得到三维数组 #

针对需求，最简便的方式是利用广播机制结合轴变换实现，效率远高于循环：

代码实现 #

import numpy as np

matrix = np.array([[10,1,32],[32,4,15],[6,3,1]])
# 广播相减：得到shape=(3,3,3)的中间结果，intermediate[i,j,k] = matrix[i,j] - matrix[i,k]
intermediate = matrix[:, None, :] - matrix[:, :, None]
# 轴变换调整维度顺序，匹配需求的输出结构
output = intermediate.transpose(2, 0, 1)

print(output)
# 输出与期望一致：
# [[[  0  -9  22]
#   [  0 -28 -17]
#   [  0  -3  -5]]
#
#  [[  9   0  31]
#   [ 28   0  11]
#   [  3   0  -2]]
#
#  [[-22 -31   0]
#   [ 17 -11   0]
#   [  5   2   0]]]

原理说明 #

• matrix[:, None, :] 将原数组扩展为(3,1,3)的三维数组，保持行维度不变，新增中间维度；
• matrix[:, :, None] 将原数组扩展为(3,3,1)的三维数组，保持行、列维度，新增末尾维度；
• 广播相减后自动匹配维度，得到每行内所有元素两两相减的结果；
• transpose(2,0,1) 将轴顺序从(行, 原列, 减数列)调整为(减数列, 行, 原列)，正好对应output[k]为每行元素减去该行第k个元素的结构。

也可以用更直观的列表推导式实现（适合小体量数组）：

output = np.array([matrix - col[:, None] for col in matrix.T])

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