Index of an element in a numpy array

Question

Given a matrix like the following:

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

How can I pinpoint the index of an element of interest. For example, assume I would like to find the index of 2 in the first row of the np.array, like so: A[0,:].index(2), but clearly this does not work because A[0,:] is not a list.

Answer 1

You can compare the array to the value 2, and then use where.

For example, to find the location of 2 in the first row of A:

In [179]: np.where(A[0, :] == 2)[0]
Out[179]: array([1])

In [180]: j = np.where(A[0, :] == 2)[0]

In [181]: A[0, j]
Out[181]: array([2])

where also works with higher-dimensional arrays. For example, to find 2 in the full array A:

In [182]: i, j = np.where(A == 2)

In [183]: A[i,j]
Out[183]: array([2])