计算 0 的数量，直到每行中的第一个 1

Question

我有一个 2D numpy 数组，我想计算值的数量（包括特定值），直到每行出现特定值。 如果该特定值不存在，则返回该行的长度。 例如：

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

我想要array([1, 5, 3])返回，因为前 2 出现在第一列的第一行； 没有出现在第二行； 并出现在第三行的第三列。

我的尝试是通过建造 2 个新的临时 arrays； 至少可以说非常笨重。

new_arr = np.zeros_like(arr)
new_arr[np.arange(len(arr)), np.argmax(np.cumsum(arr==val, 1)==1, 1)] = 1
new_new_arr = np.nonzero(new_arr==1)[1] + 1
new_new_arr[np.all(arr!=val, 1)] = arr.shape[1]

这有点像从指数分布中绘制，所以我希望 numpy 库中有一个本机方法已经实现了这样的东西，但我找不到它。 我应该如何以 numpythonic 的方式做到这一点？

Answer 1

In [180]: val = 2

In [181]: m = arr == val

In [182]: np.where(m.any(axis=1), m.argmax(axis=1) + 1, arr.shape[1])
Out[182]: array([1, 5, 3], dtype=int64)

• 从是否等于val中获取掩码
• 如果any值在一行中相等，则返回该行的argmax ，即第一个 True 的索引
  • 但加 1 因为还需要包含该值
• 否则，返回列数，即.shape[1]