基于NumPy提取二维数组前两列唯一且第三列最小的行

假设我们有一个每行含3个元素的二维NumPy数组，要实现提取前两列组合唯一的行，重复组合保留第三列最小的行，可以用以下两种高效的矢量化方法：

方法一：排序后取唯一行（推荐，性能最优） #

利用np.lexsort先按前两列分组排序，再按第三列升序排序，最后通过np.unique提取每个分组的首行（即第三列最小的行）：

import numpy as np

# 示例数组
arr = np.array([
    [1, 2, 5],
    [1, 2, 3],
    [2, 3, 7],
    [2, 3, 4],
    [3, 4, 6]
])

# 按「第一列→第二列→第三列升序」排序，确保同组内第三列最小的行排在最前
sorted_idx = np.lexsort((arr[:, 2], arr[:, 0], arr[:, 1]))
sorted_arr = arr[sorted_idx]

# 提取前两列唯一组合的首行索引，得到最终结果
unique_indices = np.unique(sorted_arr[:, :2], axis=0, return_index=True)[1]
result = sorted_arr[unique_indices]

print(result)
# 输出：
# [[1 2 3]
#  [2 3 4]
#  [3 4 6]]

方法二：分组取最小值索引 #

通过np.unique获取前两列的分组信息，再对每个分组筛选第三列最小的行：

import numpy as np

arr = np.array([
    [1, 2, 5],
    [1, 2, 3],
    [2, 3, 7],
    [2, 3, 4],
    [3, 4, 6]
])

# 获取前两列的唯一组合、分组逆序索引（标记每行属于哪个唯一组）
_, inv = np.unique(arr[:, :2], axis=0, return_inverse=True)

# 遍历每个分组，找到组内第三列最小的行索引
selected_indices = []
for group_id in np.unique(inv):
    group_rows = np.where(inv == group_id)[0]
    min_row_idx = group_rows[np.argmin(arr[group_rows, 2])]
    selected_indices.append(min_row_idx)

result = arr[selected_indices]

print(result)
# 输出同方法一

关键说明 #

• 方法一依赖NumPy底层优化的排序和去重操作，处理大数组时性能远优于手动循环，是首选方案。
• 两种方法均避免了繁琐的嵌套循环，完全利用NumPy的矢量化特性实现高效计算。

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