python - 不满足条件的第一个值的索引 Numpy

Question

我有一个类似于此的 2D numpy 数组:

np.array([
  [5, 4, 3, 2, 1], 
  [4, 3, 3, 3, 3],
  [4, 4, 2, 5, 1],
  [2, 2, 2, 2, 1] 
])

我想找到第一个违反条件的值的索引，如果它们都满足条件，则返回最大索引+ 1。例如，我想找到条件 value >=3 为 false 的第一个索引。

• 在 ls = [5,4,3,2,1] 中，它将是 3，因为value>=3 为 false 的第一个索引是 ls[3]，其值为2。

• 在ls = [4,3,3,3,3]中，由于列表中的所有值都满足value >=3，因此返回最大值索引+1，即5。

• 对于 ls= [4,4,2,5,1]，自第一个索引以来，它的值为 2违反条件的是ls[2]，它的值是2。

• 对于 ls= [2,2,2,2,1]，自第一个索引以来，它将为 0违反条件的是 ls[0]，其值为 2。

我当前尝试应用我的条件来创建 bool 掩码，然后使用 np.argmin() 查找第一个 false 值的索引

 np.apply_along_axis(
      lambda x: np.argmin(x>=3),
      1,
      np.array([
          [5, 4, 3, 2, 1], 
          [4, 3, 3, 3, 3],
          [4, 4, 3, 5, 1],
          [2, 2, 2, 2, 1] 
      ])
  )

但是这会输出array([3, 0, 2, 0])，因为当数组中的所有值都满足条件时，此方法无法处理第二种情况。有更好的方法吗？请记住，这必须应用于数千个以上的数组，因此理想情况下必须将复杂性保持在最低限度。

Answer 1

这是一种方法 -

In [60]: a
Out[60]: 
array([[5, 4, 3, 2, 1],
       [4, 3, 3, 3, 3],
       [4, 4, 2, 5, 1],
       [2, 2, 2, 2, 1]])

In [61]: m = a<3

In [62]: np.where(m.any(1),m.argmax(1),a.shape[1])
Out[62]: array([3, 5, 2, 0])