Finding indices given condition in numpy matrix

Question

So I have a numpy matrix, such as:

[[1,2,3],
[0,59,2],
[54,4,2]]

Now I want to find the indices, where all values are greater than 50 (not the maximum). This should give me [1,1],[2,0].

Other than iterating through, checking each value and keeping track of indices for values which obey condition and then returning these pairs - how do you do this efficiently?

Answer 1

You want to use either numpy.where or numpy.argwhere:

import numpy as np
A = np.array([[99, 2, 3],
              [0, 59, 2],
              [54, 4, 2]])
np.where(A > 50)
# (array([0, 1, 2]), array([0, 1, 0]))
np.argwhere(A > 50)
# array([[0, 0],
#        [1, 1],
#        [2, 0]])