Find Index of Element in Numpy Array

In this tutorial, we will look at how to find the index of an element in a numpy array.

How to find the index of element in numpy array?

You can use the numpy’s where() function to get the index of an element inside the array. The following example illustrates the usage.

np.where(arr==i)

Here, arr is the numpy array and i is the element for which you want to get the index. Inside the function, we pass arr==i which is a vectorized operation on the array arr to compare each of its elements with the value in i and result in a numpy array of boolean True and False values.

Now, np.where() gives you all the indices where the element occurs in the array. That is the indices of all the elements for which arr==i evaluates to True. This method works for both one-dimensional and multidimensional arrays. See the examples below:

1. Index of element in 1D array

Let’s apply the above syntax on a one-dimensional numpy array and find all the indices where a particular element occurs. First, let’s create a 1D array and print it out.

import numpy as np

# create a numpy array
arr = np.array([7, 5, 8, 6, 3, 9, 5, 2, 3, 5])
# print the original array
print("Original array:", arr)

Output:

Original array: [7 5 8 6 3 9 5 2 3 5]

Now that we have a 1D numpy array, let’s find the indexes where the element 5 occurs inside the array:

# find index of 5
result = np.where(arr==5)
# print the result
print("Index of 5:", result)

Output:

Index of 5: (array([1, 6, 9], dtype=int64),)

We get a tuple of numpy arrays as an output. Note that this tuple only has one numpy array storing the indices of occurrence of the element 5 inside the array.

If you were to use the np.where() function on a multidimensional numpy array, the returned tuple would have multiple numpy arrays, one for each axis.

What if the element is not present in the array?

If the element is not present in the array we get an empty array with np.where(). For example, let’s use it to find the index of 1, an element that is not present in the above array arr.

# index of 1
print(np.where(arr==1))

Output:

(array([], dtype=int64),)

You can see that the returned tuple contains an empty numpy array.

Index of the first occurrence of element

Since np.where() returns all the indexes of the occurrence of an element. You can use it to find the index of the first occurrence. For example, let’s find the index of the first occurrence of 5 in the above array.

# the first index of 5
print("First index of 5:", np.where(arr==5)[0][0])

Output:

First index of 5: 1

From the previous examples, we know that 5 is present at indexes 1, 6, and 9 in the array arr. Here we get its first occurrence which is at index 1.

2. Index of element in 2D array

We can also use the np.where() function to find the position/index of occurrences of elements in a two-dimensional or multidimensional array. For a 2D array, the returned tuple will contain two numpy arrays one for the rows and the other for the columns.

First, let’s create a two-dimensional numpy array.

# create a 2D numpy array
arr = np.array([[21, 17, 19],
                [15, 23, 17],
                [17, 11, 16]])
# print the original array
print("Original array:\n", arr)

Output:

Original array:
 [[21 17 19]
 [15 23 17]
 [17 11 16]]

Here we created a 2D numpy array with three rows and three columns. Let’s find the indexes where 17 occurs inside this array.

# find index of 17
result = np.where(arr == 17)
# print the result
print("Index of 17:", result)

Output:

Index of 17: (array([0, 1, 2], dtype=int64), array([1, 2, 0], dtype=int64))

The returned tuple from np.where() contains two numpy arrays. The first array values tell the row indexes whereas the second array values tell the column indexes of the occurrences of the element inside the 2D array.

Let’s make these indexes more readable by showing the (row, column) index tuples for each occurrence.

# get the position in the 2-d array
print(list(zip(result[0], result[1])))

Output:

[(0, 1), (1, 2), (2, 0)]

The result shows that 17 occurs at the following locations – row 0 column 1, row 1 column 2, and row 2 column 0 with the index for rows and columns starting from 0.

What if the element is not present in the array?

If the element is not present in the 2D array. The returned tuple from np.where() will have two empty numpy arrays. For example, if we check for the index of 13, an element that is not present in the 2D array above –

print("Index of 13:", np.where(arr==13))

Output:

Index of 13: (array([], dtype=int64), array([], dtype=int64))

We get a tuple of two empty numpy arrays.

For more on the numpy where function, refer to its documentation.

With this, we come to the end of this tutorial. The code examples and results presented in this tutorial have been implemented in a Jupyter Notebook with a python (version 3.8.3) kernel having numpy version 1.18.5