Numpy way of selecting indices of 2D numpy array if a condition met [closed]

Question

I have a snippet of code to get indices of a 2D numpy array if the maximum value of that raw is greater than a number.

index = []
for i in range(np.shape(soft_p)[0]):
      if soft_p[i].max() > 0.4:
           index.append(i)

I was wondering if there's a better way to do it using numpy conditions and functions

Answer 1

You should find the maximum value according to the first axis (row), and then select the row number greater than 0.4:

max_vals = soft_p.max(1)
index = (max_vals > 0.4).nonzero()[0]

Example:

>>> ar = np.random.rand(5, 5)
>>> ar
array([[0.46621005, 0.51573283, 0.53287116, 0.82355335, 0.95846324],
       [0.76834591, 0.63605015, 0.32834301, 0.15062014, 0.69269814],
       [0.94558026, 0.11578385, 0.43292249, 0.08532496, 0.16731557],
       [0.14552061, 0.87651804, 0.25594289, 0.22702692, 0.23300077],
       [0.19116528, 0.54291561, 0.49146977, 0.98988884, 0.43468107]])
>>> max_vals = ar.max(1)
>>> max_vals
array([0.95846324, 0.76834591, 0.94558026, 0.87651804, 0.98988884])
>>> (max_vals > 0.9).nonzero()[0]
array([0, 2, 4], dtype=int64)