在Python中基于数组索引执行元素差值计算的实现方法

问题需求 #

我有两个Numpy数组I和X，需要完成这样的操作：提取数组I中的索引对，用X对应索引的元素计算差值。比如I[0]是[0,1]，就计算X[0]-X[1]并加入新数组T；I[1]是[1,2]，就计算X[1]-X[2]加入T。

初始代码示例 #

import numpy as np
I = np.array([[0,1],[1,2]])
X = np.array([10,5,3])

期望输出 #

T = np.array([[5], [2]])
# 对应计算逻辑：10-5=5，5-3=2

解决方案 #

下面提供几种简洁高效的Numpy实现方式：

方法一：直接索引差值计算 #

利用Numpy的数组索引特性，直接提取对应位置元素做差：

import numpy as np
I = np.array([[0,1],[1,2]])
X = np.array([10,5,3])

# 提取I每行第一个索引的X元素，减去每行第二个索引的X元素，再调整为列数组
T = (X[I[:, 0]] - X[I[:, 1]]).reshape(-1, 1)
print(T)

运行输出：

[[5]
 [2]]

方法二：批量提取后运算 #

先一次性取出所有需要的元素对，再逐行计算差值：

x_pairs = X[I]  # 得到形状为(2,2)的数组：[[10,5],[5,3]]
T = (x_pairs[:, 0] - x_pairs[:, 1]).reshape(-1, 1)
print(T)

方法三：加权矩阵运算（扩展性更强） #

如果后续需要支持更复杂的加权运算，这种方式更通用：

x_pairs = X[I]
weights = np.array([1, -1])  # 差值等价于第一个元素乘1加第二个元素乘-1
T = (x_pairs @ weights).reshape(-1, 1)
print(T)

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