Get row numbers of rows matching a condition in numpy

Question

Suppose I have a numpy array like:

a = array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9],
       [3, 2, 1]])

I want to check if the second element == 2.

I know I can do this:

>>> a[:,1]==2
array([ True, False, False,  True], dtype=bool)

returning booleans. My question is, how do I get the row numbers of the rows where the condition is true? In this example I would want to get back array([0, 3]) because the 0th and 3rd rows match the condition second element == 2.

Answer 1

This is the function I used, which provides only the row numbers:

import numpy as np

a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [3, 2, 1]])

row_locator = a[:,1]==2

row_values = []

for x in range(0, len(a)):
    if row_locator[x]:
        row_values.append(x)

print(row_values)

>>>[0, 3]

Answer 2

Use np.where to return the indices:

In [79]:

np.where(a[:,1]==2)
Out[79]:
(array([0, 3], dtype=int64),)