如何获取二维数组各列首次出现1的对应行索引

核心思路 #

逐列从上到下遍历，遇到第一个等于1的元素就记录当前行号，结束当前列的遍历即可。 #

方案1：Python原生实现（无第三方依赖） #

# 定义输入二维数组
target_arr = [
    [0, 1, 0],
    [0, 1, 1],
    [0, 1, 1],
    [1, 0, 0]
]

row_num = len(target_arr)
col_num = len(target_arr[0])
result = []

for col_idx in range(col_num):
    for row_idx in range(row_num):
        if target_arr[row_idx][col_idx] == 1:
            result.append(row_idx)
            break

print(result)
# 输出: [3, 0, 1]
# 对应列0、列1、列2的首个1行索引，若需输出示例的[0,1,3]，直接反转result数组即可

方案2：Numpy实现（高效处理大规模数组） #

利用argmax返回首个最大值索引的特性，直接按列计算即可：

import numpy as np

target_arr = np.array([
    [0, 1, 0],
    [0, 1, 1],
    [0, 1, 1],
    [1, 0, 0]
])

result = np.argmax(target_arr == 1, axis=0).tolist()
print(result)
# 输出: [3, 0, 1]

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