Index of First Value Where Condition Not Met Numpy

Question

I have a 2D numpy array that looks akin to this:

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

I want to find the index of the first value that violates a condition, and if they all satisfy the condition, return the maximum index + 1. For example, I want to find the first index where the condition value >=3 is false.

In ls = [5,4,3,2,1] it would be 3, as the first index where value>=3 is false is ls[3] which has the value 2.
In ls = [4,3,3,3,3], since all values in the list satisfy value >=3, return the max index +1, which is 5.
For ls= [4,4,2,5,1] it would therefore be 2 since the first index that violates the condition is ls[2] which is the value 2.
For ls= [2,2,2,2,1] it would therefore be 0 since the first index that violates the condition is ls[0] which is the value 2.

My current attempt applies my condition to create a boolean mask, then uses np.argmin() to find the index of the first false value

 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] 
      ])
  )

But this outputs array([3, 0, 2, 0]) since this method cannot handle the second case when all values in the array satisfy the condition. Is there a better approach? Please keep in mind this has to be applied to over a few thousand of these arrays so ideally complexity must be kept to a minimum.

Divakar · Accepted Answer · 2020-09-01 10:37:21Z

5

Here's one way -

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])

answered Sep 1, 2020 at 10:37

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Index of First Value Where Condition Not Met Numpy

1 Answer 1

Comments

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Related