Advanced Indexing

NumPy supports different indexing techniques for accessing subarrays. We already discussed list-like indexing. Now we add boolean and integer indexing.

import numpy as np

Boolean Indexing

If the index to an array is a boolean array of the same shape as the indexed array, then a one-dimensional array of all items where the index is True is returned.

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

b = a[idx]

print(a, '\n')
print(idx, '\n')
print(b)

[[1 2 3]
 [4 5 6]
 [7 8 9]] 

[[ True  True False]
 [False  True  True]
 [ True False False]] 

[1 2 5 6 7]

A typical use case are elementwise bounds:

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

b = a[a > 3]

print(b)

[4 5 7 6 4 5 6 7 4 9]

Here b is an array containing all numbers of a which are greater than 3. The comparison a > 3 returns a boolean array of the same shape as a. Note, that broadcasting is used to compare an array to a number. The resulting boolean array then is used as index to a.

Integer Indexing

Given an array we may provide an array of indices. The result of the corresponding indexing operation is an array of the same size as the index array, but with the items of the indexed array at corresponding positions.

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

print(a[idx])

[[1 6]
 [3 4]]

For indexing multi-dimensional arrays we need multiple index arrays (one per dimension).

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

print(a[idx0, idx1])

[[1 3]
 [8 7]]