Where in 2D

This follows from where and argmin.

That page covers how np.where and np.argmin work for one-dimensional arrays.

This page covers what happens for np.where with arrays of two dimensions and more.

Arrays can have two dimensions

import numpy as np

So far we have only seen one-dimensional arrays.

A vector is another term for a one-dimensional array.

Here’s a one-dimensional array (vector):

arr_one_d = np.array([2, 99, -1, 4, 99])
arr_one_d
array([ 2, 99, -1,  4, 99])

Notice the shape of the array:

arr_one_d.shape
(5,)

You can also think of this one-dimensional array as a single row.

In fact you can have arrays with more than one row, like this:

arr_two_d = np.array([[1, 2, 3, -2, -3], [10, 11, 12, -3, -13]])
arr_two_d
array([[  1,   2,   3,  -2,  -3],
       [ 10,  11,  12,  -3, -13]])

Notice the shape. This array has 2 rows, each with 5 columns.

arr_two_d.shape
(2, 5)

Remember, we can ask for a single element from a one-dimensional array, using indexing with an integer between square brackets. The integer gives the position (or offset) of the element we want.

arr_one_d[1]
99

With a two-dimensional array, we need to specify the row and column of the element we want:

arr_two_d[1, 4]
-13

The first value between the brackets is the row position, and the second is the column position.

Revision - using where on one-dimensional arrays

np.where gets the indices of the True values in a Boolean array.

# A Boolean array.
equal_to_99 = arr_one_d == 99
equal_to_99
array([False,  True, False, False,  True])
# Indices of the True values.
indices = np.where(equal_to_99)
indices
(array([1, 4]),)

We can use the returned indices to index into the array, using square brackets.

arr_one_d[indices]
array([99, 99])

The two-dimensional case

This also works in two or more dimensions.

np.where now returns two index arrays, one for the rows, and one for the columns.

indices2d = np.where(arr_two_d == -3)
indices2d
(array([0, 1]), array([4, 3]))

Just as for the one-dimensional case, we can use the returned indices to index into the array, and get the elements.

arr_two_d[indices2d]
array([-3, -3])