如何获取numpy数组对角区域中值1首次出现的列？

问题场景 #

我们有如下NumPy数组：

import numpy as np
b = np.array([[0.3, -0.2, 0.4, 0.5, -0.8, 1.0, 0.0, 0.0],
             [0.6, 0.2, 0.7, 0.91, 0.67, 0.0, 1.0, 0.0],
             [0.5, 0.1, 0.7, 0.0, 0.6, 0.0, 0.0, 1.0]])

数组右侧最后3列为对角矩阵，我们需要获取其中值1首次出现的列（索引为5）。使用代码first_occurence = np.argmax(b == 1, axis=1)[0]能得到正确结果，但对于下面的数组：

b = np.array([[0.3, -0.2, 0.4, 0.5, -0.8, 0.0, 0.0, 0.0],
              [0.6, 0.2, 0.7, 0.91, 0.67, 0.0, 1.0, 0.0],
              [0.5, 0.1, 0.7, 0.0, 0.6, 0.0, 0.0, 1.0]])

上述方法返回0，而正确结果应为索引6，该如何解决？

解决方法 #

方案1：过滤无效行后取索引 #

np.argmax在输入数组全为False时会返回0，这是导致错误的核心原因。我们可以先筛选出包含1的行，再对这些行取首次出现的索引：

# 定位所有存在1的行索引
valid_rows = np.where(np.any(b == 1, axis=1))[0]
# 取第一个有效行中1首次出现的列索引
first_occurrence = np.argmax(b == 1, axis=1)[valid_rows[0]]

方案2：直接用非零索引定位 #

利用np.nonzero直接获取所有值为1的位置，再取第一个结果的列索引，逻辑更简洁：

# 获取所有值为1的位置的行列索引
rows, cols = np.nonzero(b == 1)
# 取第一个出现的列索引
first_occurrence = cols[0]

方案3：处理无1的边界情况 #

如果需要兼容数组中完全没有1的场景，可以增加判断逻辑：

if np.any(b == 1):
    rows, cols = np.nonzero(b == 1)
    first_occurrence = cols[0]
else:
    # 自定义未找到时的返回值，比如-1
    first_occurrence = -1

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