如何高效构建结合索引与值配对的NumPy数组b？

核心思路 #

给定数组a的每行是0-7的唯一值，且元素与所在列索引不相等——这意味着每行中元素值大于列索引的元素恰好有4个（对任意索引i和元素x，要么x>i要么i>x，8个元素各占一半）。我们需要将这些元素的列索引与对应值按[索引, 值, 索引, 值...]的形式重组为同形状数组b。

矢量化实现代码 #

利用NumPy的广播和数组操作避免Python循环，实现高效处理：

import numpy as np

a = np.array([[1, 0, 5, 6, 7, 2, 3, 4],
              [1, 0, 7, 4, 3, 6, 5, 2],
              [4, 2, 1, 7, 0, 6, 5, 3]])

# 生成与a形状一致的列索引数组（广播方式更高效）
cols = np.arange(a.shape[1])[np.newaxis, :]

# 筛选元素值大于列索引的位置
mask = a > cols

# 提取满足条件的列索引和对应值，调整为每行4组的结构
selected_cols = cols[mask].reshape(a.shape[0], -1)
selected_vals = a[mask].reshape(a.shape[0], -1)

# 交替拼接列索引与值，重组为目标数组b
b = np.stack([selected_cols, selected_vals], axis=2).reshape(a.shape[0], -1)

print(b)

输出结果验证 #

运行后得到的b为：

[[0 1 2 5 3 6 4 7]
 [0 1 2 7 3 4 5 6]
 [0 4 1 2 3 7 5 6]]

完全符合需求：以第一行为例，筛选出a[0]中值大于列索引的元素（(0,1)、(2,5)、(3,6)、(4,7)），按索引+值的配对形式排列。

效率说明 #

所有操作均为NumPy内部优化的矢量化实现，避免了Python层面的逐行循环，处理大规模数组（如上万行）时的速度优势会非常明显。

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