如何获取NumPy二维数组每行最后一个1的索引？

针对给定的左侧连续为1的二维NumPy数组，这里提供几种高效的解决方法：

方法一：反转数组后取第一个1的镜像索引 #

利用数组每行1左连续的特性，反转每行后找到第一个1的位置，再通过计算得到原数组中最后一个1的索引：

import numpy as np

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

# 反转每行元素
reversed_arr = arr[:, ::-1]
# 计算原数组中最后一个1的索引
last_one_idx = reversed_arr.shape[1] - 1 - reversed_arr.argmax(axis=1)

print(last_one_idx)
# 输出：array([2, 0, 1, 3, 0])

方法二：利用累计求和找最大值索引 #

对每行从左到右累计求和，累计和的最大值位置即为最后一个1的索引（后续0不会改变累计和，argmax返回的第一个最大值位置正好对应最后一个1的位置）：

cum_sum = arr.cumsum(axis=1)
last_one_idx = cum_sum.argmax(axis=1)

print(last_one_idx)
# 输出：array([2, 0, 1, 3, 0])

方法三：遍历每行查找最后一个1（适合小数据量） #

通过列表推导式遍历每行，用np.where定位所有1的索引后取最后一个：

last_one_idx = np.array([np.where(row == 1)[0][-1] for row in arr])

print(last_one_idx)
# 输出：array([2, 0, 1, 3, 0])

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