計算 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]