如何用NumPy/SciPy高效实现二维数组列间自差分（替代循环）

我有一个形状约为(3000, 5)的二维NumPy数组mat，需要计算任意两列间的差值，但只保留i<j时第i列减第j列的结果（忽略自减和j<i的情况）。目前已经用嵌套for循环实现了功能，但想借助NumPy或SciPy的内置函数简化步骤，提升运行效率。

当前使用的循环实现代码：

import numpy as np

agreement = np.full((mat.shape[1], mat.shape[1]), [np.nan]*mat.shape[0])
for i in range(agreement.shape[0] - 1):
    for j in range(i+1, agreement.shape[1]):
        A = mat[:, i]
        B = mat[:, j]
        diff = A - B
        agreement[i][j] = diff

最终得到的agreement数组结构示例：

array([[nan [......] [......] [......] [......]]
       [nan nan [......] [......] [......]]
       [nan nan nan [......] [......]]
       [nan nan nan nan [......]]
       [nan nan nan nan nan]])

其中[......]代表对应两列的差值结果（形状为(3000,)），nan是由[np.nan]*mat.shape[0]组成的数组。

利用NumPy的向量化操作和上三角掩码可以一次性完成计算，彻底避免循环，大幅提升效率：

简洁实现代码 #

import numpy as np

n_cols = mat.shape[1]
# 一次性计算所有列对的差值，维度为(3000, n_cols, n_cols)
all_col_diffs = mat[:, :, np.newaxis] - mat[:, np.newaxis, :]
# 生成上三角掩码（仅i<j的位置为True），扩展维度后匹配差值数组
upper_mask = np.triu(np.ones((n_cols, n_cols), dtype=bool), k=1)[..., np.newaxis]
# 用掩码筛选出需要保留的差值，其余位置填充为nan
agreement = np.where(upper_mask, all_col_diffs.transpose(1, 2, 0), np.nan)

代码解释 #

1. 广播相减：通过给mat添加新维度，让mat[:, :, np.newaxis]（形状(3000, 5, 1)）和mat[:, np.newaxis, :]（形状(3000, 1, 5)）触发广播机制，一次性计算出所有列对的差值，结果all_col_diffs[:, i, j]就是第i列减第j列的结果。
2. 上三角掩码：np.triu(..., k=1)生成一个上三角布尔矩阵，对角线及以下位置为False，仅i<j的位置为True。通过[..., np.newaxis]扩展维度后，和all_col_diffs.transpose(1,2,0)（形状(5,5,3000)）的维度完全匹配。
3. 赋值与填充：np.where根据掩码筛选出需要保留的差值，其余位置自动填充为nan，一步得到目标结构的agreement数组。

这种向量化实现的效率远高于循环，当数组列数越多时，性能提升越明显。

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