如何用Numpy高效按首行唯一值排序并提取对应第二行最大值？

针对大数据量的场景，完全可以用Numpy的向量式操作替代循环，以下是高效解决方案：

步骤说明 #

1. 对首行元素排序，同步调整第二行的对应位置，保证两组数据的对应关系
2. 定位排序后首行中唯一值的起始索引，以此作为分组边界
3. 利用np.maximum.reduceat批量计算每个分组内第二行的最大值
4. 整理排序后的唯一值和对应最大值为目标数组

代码实现 #

import numpy as np

a = np.array([[1,2,1,2],[5,3,3,8]])

# 提取两行数据
row1, row2 = a[0], a[1]

# 按首行元素排序，获取排序索引
sorted_indices = np.argsort(row1)
sorted_row1 = row1[sorted_indices]
sorted_row2 = row2[sorted_indices]

# 获取排序后首行的唯一值及其起始索引
unique_vals, group_starts = np.unique(sorted_row1, return_index=True)

# 按分组计算第二行的最大值
group_max = np.maximum.reduceat(sorted_row2, group_starts)

# 组装结果数组
result = np.array([unique_vals, group_max])
print(result)
# 输出: [[1 2]
#        [5 8]]

效率优势 #

所有操作均为Numpy底层优化的向量运算，完全避免Python层面的循环，能够高效处理超大维度的数组——相比逐元素循环，性能提升可达数十甚至上百倍。

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