本文介绍如何在不使用显式循环的前提下，对二维 NumPy 数组按第三行（索引列）分组，并高效提取每组中第一行（数值列）的最大值所对应的整行数据。核心方法是利用 np.lexsort 进行多关键字排序与布尔索引组合。
本文介绍如何在不使用显式循环的前提下，对二维 numpy 数组按第三行（索引列）分组，并高效提取每组中第一行（数值列）的最大值所对应的整行数据。核心方法是利用 `np.lexsort` 进行多关键字排序与布尔索引组合。
在科学计算和数据预处理中，常需对结构化二维数组（如 [values, aux_data, indices]）按索引列分组，并提取每组中某主键（如 values）最大时对应的完整记录。若直接使用 np.maximum.reduceat，它仅支持沿一维轴聚合标量结果（如仅返回最大值），无法同步保留其他列的关联数据——这正是原始问题的瓶颈。
此时，np.lexsort 成为更优雅的解决方案：它支持多列联合排序，且稳定、向量化、无 Python 循环。其关键思想是先按索引列升序、再按数值列降序排序，这样每个索引组内的首元素即为该组最大值所在行；随后通过检测索引列的“组边界”（即 data[2, i] != data[2, i+1]），即可一次性提取所有目标行。
以下为完整实现：
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. lexsort 排序键：先按索引列 (data[2]) 升序，再按数值列 (data[0]) 降序（负号实现）
# → 确保每组内最大值排在最前
sort_indices = np.lexsort((-data[0], data[2]))
# 2. 按排序索引重排 data
sorted_data = data[:, sort_indices]
# 3. 找出每个新索引组的首个位置（即组切换点）
# 使用布尔掩码：当前索引 ≠ 下一个索引，且末尾强制为 True
group_boundaries = np.r_[sorted_data[2, :-1] != sorted_data[2, 1:], True]
# 4. 提取每组首行（即最大值所在行）的所有列
result = sorted_data[:, group_boundaries]
print(result)
输出：
[[ 0.66176217 0.45114132] [10. 4. ] [ 0. 1. ]]
✅ 优势总结：
- 完全向量化：零 Python 循环，底层由 NumPy C 实现，适合百万级数据；
- 内存友好：仅需一次排序与一次布尔索引，无需中间列表或字典；
- 可扩展性强：若需最小值，将 -data[0] 改为 data[0] 即可；若需多列辅助数据，data 形状可任意扩展（只要索引列固定为某一行）；
- 稳定性保障：lexsort 是稳定排序，相同索引内顺序由次关键字严格定义。
⚠️ 注意事项：
- 索引列（本例中 data[2]）必须为可比较类型（如 int/float），且建议提前去除非有限值（np.inf, np.nan）；
- 若索引列含浮点数，注意精度问题，必要时用 np.round(data[2]).astype(int) 转换；
- np.r_[] 是拼接工具，[:-1] != [1:] 生成长度为 n-1 的布尔数组，末尾 True 补齐至 n 长度，确保最后一组被包含。
该方法在性能与可读性间取得良好平衡，是 NumPy 高级分组聚合任务的标准实践之一。
