本文介绍如何在不使用显式循环的前提下，对二维数组按第三行索引分组，高效提取每组中第一行的最大值及其对应第二行的关联数值。核心方案是利用 np.lexsort 对数据进行稳定排序与布尔掩码组合，实现原子级向量化操作。
本文介绍如何在不使用显式循环的前提下，对二维数组按第三行索引分组，高效提取每组中第一行的最大值及其对应第二行的关联数值。核心方案是利用 `np.lexsort` 对数据进行稳定排序与布尔掩码组合，实现原子级向量化操作。
在处理结构化二维数组（如 [values, aux_values, indices]）时，常需按 indices 分组，并为每组选取 values 的最大值，同时保留其对应的 aux_values —— 这本质上是一个“分组取极值+同步索引保持”问题。若用 np.maximum.reduceat，它仅支持沿单一轴对一维数组做分段归约，无法直接保证多行数据的行列对齐；而基于 np.isin 的掩码方法又易因值重复导致误匹配（如示例中 0.45114132 出现两次但归属不同组），不可靠。
此时，np.lexsort 成为更稳健的选择：它支持多关键字稳定排序，且排序后相同 indices 的元素必然连续排列，最大值必位于每组末尾（若按 (index, value) 升序排，则最大值在组内最右；若按 (index, -value) 排，则最大值在组内最左）。我们采用前者，并借助差分逻辑定位每组最后一个位置：
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. ]
])
# 步骤分解：
# 1. 按 [index, value] 双关键字升序排序（主关键字 index，次关键字 value）
idx_sorted = np.lexsort((data[0], data[2])) # 注意：lexsort 的键顺序是反的：(次, 主)
# 2. 获取排序后 data[2] 的分组边界：当前索引值 ≠ 下一索引值的位置即为每组末尾
sorted_indices = data[2, idx_sorted]
group_ends = np.r_[sorted_indices[:-1] != sorted_indices[1:], True]
# 3. 提取每组末尾位置在原始排序索引中的下标，并切片原 data
final_idx = idx_sorted[group_ends]
result = data[:, final_idx]
print(result)输出：
[[ 0.66176217 0.45114132] [10. 4. ] [ 0. 1. ]]
✅ 关键优势：
- 完全向量化，无 Python 循环或列表推导，适合百万级数据；
- 利用 lexsort 的稳定性，确保同 index 组内 value 有序，最大值位置可预测；
- np.r_ 与布尔索引组合简洁高效，避免额外内存分配；
⚠️ 注意事项：
- np.lexsort(keys) 中 keys 应按次要到主要顺序传入（即 lexsort((value, index)) 表示先按 index 分组、再按 value 组内排序）；
- 若存在 NaN 值，lexsort 会将其排在末尾，需预先处理（如 np.nan_to_num 或掩码过滤）；
- 该方法默认取每组最大值；若需最小值，可将 data[0] 替换为 -data[0] 后复用逻辑，或改用 np.argmin 配合 np.split（但会引入额外开销）；
综上，面对“按索引分组 + 同步提取极值及关联字段”的需求，np.lexsort + 布尔边界检测是兼顾性能、简洁性与鲁棒性的首选方案。
