本文介绍如何在不使用显式循环的前提下，对二维 NumPy 数组按第三行（索引列）分组，并高效提取每组中第一行最大值所对应的整行数据。核心方法是利用 np.lexsort 对数组进行多关键字排序后做逻辑切片，兼顾性能与简洁性。
本文介绍如何在不使用显式循环的前提下，对二维 numpy 数组按第三行（索引列）分组，并高效提取每组中第一行最大值所对应的整行数据。核心方法是利用 `np.lexsort` 对数组进行多关键字排序后做逻辑切片，兼顾性能与简洁性。
在处理结构化数值数据时，常遇到“按类别索引聚合并保留关联字段”的需求：例如一个形状为 (3, k) 的数组，其中第 0 行为待比较的浮点值（如得分），第 1 行为需同步保留的关联数据（如 ID），第 2 行为离散索引（如类别标签）。目标是：对每个唯一索引，找出其对应位置中第 0 行的最大值，并返回该最大值所在位置的整列数据（即第 0、1、2 行均取同一列）。
直接使用 np.maximum.reduceat 只能作用于一维数组，无法天然保持多行间列对齐关系；而基于布尔掩码（如 np.isin）的方法易因值重复导致误匹配（如示例中 0.45114132 出现两次但属于不同索引组），不可靠。
推荐解法是借助 np.lexsort——它支持按多列优先级排序，且返回的是索引数组，可精准控制原始数据的重排顺序。关键思路如下：
- 以索引列（第 2 行）为主序、值列（第 0 行）为次序升序排序 → 这样每组内最大值会落在该组末尾；
- 识别每组末尾位置：对排序后的索引列做差分，检测索引变化处（配合 np.r_ 补 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. ]
])
# 步骤分解（便于理解）：
idx_sorted = np.lexsort((data[0], data[2])) # 先按 data[2] 升序，同索引内按 data[0] 升序
# → 索引序列确保每组内最大值在最后
# 标记每组最后一个元素的位置（即索引变化前一位 + 最后一个）
group_ends = np.r_[data[2, :-1] != data[2, 1:], True]
# → [True, False, True, False, True] → 对应索引 [0,2,4]
# 提取排序后数组中 group_ends 为 True 的列，并还原为原始 data 形状
result = data[:, idx_sorted[group_ends]]
print(result)
# 输出：
# [[ 0.66176217 0.45114132]
# [10. 4. ]
# [ 0. 1. ]]
✅ 优势总结：
- 零循环、纯向量化：全程使用 NumPy 原生操作，适用于百万级数据；
- 严格保序保列：lexsort 返回索引而非副本，data[:, idx] 确保三行数据列对齐无错位；
- 鲁棒性强：不依赖值唯一性，即使 data[0] 存在重复值或 data[2] 非连续整数（如 [0,0,2,2,2]）均可正确分组；
- 内存友好：仅生成必要中间布尔数组和索引数组，无冗余拷贝。
⚠️ 注意事项：
- np.lexsort(keys) 中 keys 传入顺序为次要关键字到主要关键字（与 sorted(..., key=lambda x: (x[1], x[0])) 逻辑一致），此处 (data[0], data[2]) 表示「先按 data[2] 分组，组内再按 data[0] 排序」；
- 若需取最小值而非最大值，将 data[0] 替换为 -data[0] 即可（避免改变原始数据符号）；
- 当 data[2] 为空或全相同，group_ends 仍能正确生成单元素布尔掩码，行为可预测。
该方法是 NumPy 高级索引技巧的典型应用，在科学计算与数据预处理中具有广泛适用性。
