如何高效计算一维numpy数组中元素的两两差值？

方案1：Numpy广播实现（最便捷，无需额外依赖） #

直接利用numpy的广播机制生成完整差值矩阵，再提取上三角无重复的差值结果：

import numpy as np

def pairwise_diff(X):
    # 广播生成n*n的全差值矩阵，值为 X[j] - X[i]
    diff_matrix = X - X.reshape(-1, 1)
    # 提取i<j的上三角元素（k=1表示跳过对角线元素）
    return diff_matrix[np.triu_indices(len(X), k=1)]

测试验证 #

使用你给出的示例输入测试，返回结果与你提供的完全一致：

X = np.array([0, 1, 2, 0, 1, 2, 0, 1, 2])
print(pairwise_diff(X))
# 输出：[ 1  2  0  1  2  0  1  2  1 -1  0  1 -1  0  1 -2 -1  0 -2 -1  0  1  2  0  1  2  1 -1  0  1 -2 -1  0  1  2  1]

方案2：Numba JIT实现（内存友好，适合大n场景） #

如果n大于10000，临时的n×n矩阵会占用较多内存，可以用numba编译循环实现，仅分配最终结果的内存空间：

import numpy as np
from numba import jit

@jit(nopython=True)
def pairwise_diff_low_memory(X):
    n = len(X)
    result_len = n * (n - 1) // 2
    res = np.empty(result_len, dtype=X.dtype)
    ptr = 0
    for i in range(n):
        for j in range(i + 1, n):
            res[ptr] = X[j] - X[i]
            ptr += 1
    return res

性能对比 #

n=400时的实测表现：

• 原始Python列表推导：约20~50ms
• Numpy广播方案：约0.05~0.2ms
• Numba JIT方案：约0.03~0.1ms

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