使用 np.lexsort 实现按索引分组的二维数组最大值提取
本文介绍如何在不使用显式循环的前提下，对二维 NumPy 数组按第三行（索引列）分组，并高效提取每组中第一行的最大值及其对应第二行的关联值。核心方法是利用 np.lexsort 对数据进行稳定排序与布尔掩码筛选。
本文介绍如何在不使用显式循环的前提下，对二维 numpy 数组按第三行（索引列）分组，并高效提取每组中第一行的最大值及其对应第二行的关联值。核心方法是利用 `np.lexsort` 对数据进行稳定排序与布尔掩码筛选。
在处理结构化数值数据时，常需按某列（如类别索引）分组，并提取每组中另一列（如置信度、得分）的最大值，同时保留其关联字段（如 ID、时间戳等）。对于形状为 (3, k) 的二维数组 data，其中 data[0] 为待比较数值、data[2] 为分组索引、data[1] 为需同步提取的关联值，目标是：对每个唯一索引，返回该索引下 data[0] 的最大值，以及同一位置的 data[1] 和 data[2] 值。
直接使用 np.maximum.reduceat 仅适用于一维聚合，无法天然保持多行间的列对齐关系。而基于 np.unique(..., return_index=True) 的方案虽能获取首个出现位置，但 reduceat 本身不提供“取最大值所在原始索引”的能力，导致无法准确拉取 data[1] 和 data[2] 的对应元素。
此时，np.lexsort 成为更优解：它支持多关键字稳定排序（主键优先级由参数顺序决定），且排序后相同索引的元素将连续排列，最大值自然位于每组末尾（若按升序排）或开头（若按降序排）。我们采用“先按索引 data[2] 升序，再按数值 data[0] 升序”双关键字排序，确保每组内数值递增——这样每组最后一个元素即为该索引下的最大值。
关键技巧在于构造布尔掩码以提取每组末尾元素：
故事AI绘图神器
文本生成图文视频的AI工具，无需配音，无需剪辑，快速成片，角色固定。
下载
- np.lexsort(data[[0, 2]]) 返回按 data[2]（主）、data[0]（次）升序排列的索引；
- 排序后，相同 data[2] 的元素连续，只需检测 data[2] 的相邻变化点；
- data[2,:-1] != data[2,1:] 标记组边界（前一个索引 ≠ 后一个索引的位置），再拼接 True（保留最后一组末尾）得到每组末尾索引掩码。
完整实现如下：
import numpy as np
data = np.array([
[0.45114132, 0.31522008, 0.66176217, 0.45114132, 0.26872137],
[11. , 6. , 10. , 4. , 8. ],
[0. , 0. , 0. , 1. , 1. ]
])
# 按索引(data[2])主序、数值(data[0])次序升序排序
sorted_idx = np.lexsort(data[[0, 2]]) # 注意：lexsort 以最后维度为第一关键字，故传入 [0,2] 表示先按 data[2] 再按 data[0]
sorted_data = data[:, sorted_idx]
# 构造掩码：标记每组（相同 data[2]）的最后一个位置
mask = np.r_[sorted_data[2, :-1] != sorted_data[2, 1:], True]
# 提取结果
result = sorted_data[:, mask]
print(result)
# 输出:
# [[0.66176217 0.45114132]
# [10. 4. ]
# [ 0. 1. ]]
⚠️ 注意事项：
- np.lexsort(keys) 中 keys 是一个序列，越靠后的数组优先级越高。因此 data[[0, 2]] 表示先按 data[2]（索引）分组，组内再按 data[0]（数值）排序；
- 若 data[2] 含非整数或浮点索引（如 0.0, 1.0），需确保其相等性判断稳定（浮点精度问题下建议先 np.round() 或转为整型）；
- 该方法时间复杂度为 O(k log k)，远优于 Python 循环，且完全向量化，适合百万级数据；
- 若需降序取最大值（即每组首元素），可改用 np.lexsort([-data[0], data[2]]) 并用 mask = np.r_[True, sorted_data[2, 1:] != sorted_data[2, :-1]]。
综上，np.lexsort + 掩码筛选是解决“分组取极值并保关联字段”问题的简洁、高效、纯 NumPy 方案，兼顾可读性与性能，是科学计算中值得掌握的核心技巧之一。