如何在NumPy数组中选取索引差固定的元素对？

先看你的示例数组生成代码：

import numpy as np
np.random.seed(0)
data = np.random.normal(size=(5,5))

要实现沿每行选取索引差为指定值的元素对，完全可以用NumPy的矢量化操作替代循环，效率提升明显，具体实现如下：

核心代码 #

k = 4  # 你指定的索引差
# 提取两组元素：第一组是每行中所有可作为起始点的元素，第二组是对应索引+k的元素
part1 = data[:, :-k]
part2 = data[:, k:]
# 将两组元素按行配对，得到最终结果
res = np.hstack([part1.reshape(-1,1), part2.reshape(-1,1)])

代码解释 #

• data[:, :-k]：取每行中前总列数 - k个元素，这些元素的索引i满足i + k < 总列数，是每对的第一个元素；
• data[:, k:]：取每行中从索引k开始到末尾的元素，是每对的第二个元素；
• 最后通过reshape和hstack把两组元素按顺序配对，得到(N,2)格式的结果数组，其中N是所有符合条件的元素对总数。

比如当k=4时，你的5列数组每行只有1对有效元素，最终res会是5行2列的数组，对应每行索引0和4的元素；如果k=1，每行有4对元素，最终res会是20行2列的数组，覆盖所有相邻元素对。

这种方式完全依托NumPy底层的C实现操作，没有Python层面的循环，在处理大数组时性能优势会非常显著。

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