Python中查找二维布尔数组每列最上方True值索引的高效方法

核心实现思路 #

利用numpy向量化操作避免Python层面循环，np.argmax沿列轴计算时会返回第一个最大值的下标，布尔数组中True等价于1、False等价于0，因此argmax(axis=0)会直接返回每列最靠上的True对应的行索引，再过滤掉全False的列即可得到结果。

代码实现 #

import numpy as np

# 示例布尔数组
bool_arr = np.array([
    [False, False, False, False, False, False],
    [True, False, False, True, False, False],
    [True, False, True, False, False, True],
    [True, False, True, True, False, False]
], dtype=bool)

# 计算每列第一个True的行索引
row_idx = bool_arr.argmax(axis=0)
# 生成对应列索引
col_idx = np.arange(bool_arr.shape[1])
# 过滤全为False的列：全False列argmax返回0，但对应位置值为False
valid_mask = bool_arr[row_idx, col_idx]
# 组合为最终索引元组列表
res = list(zip(row_idx[valid_mask], col_idx[valid_mask]))

运行后res输出为[(1, 0), (2, 2), (1, 3), (2, 5)]，完全符合预期。

效率说明 #

该方案全程为numpy底层C实现的向量化运算，无Python遍历开销，性能远高于np.where后逐列处理、或skyline算法实现，在百万级以上元素的大数组场景下，性能提升可达数十到上百倍。

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