NumPy - Advanced Indexing

Beyond basic slicing (e.g., arr[start:stop:step]), NumPy offers more sophisticated ways to select elements from arrays, collectively known as Advanced Indexing. This typically involves using arrays of indices (integers) or boolean arrays (masks) to access array elements.

A key difference from basic slicing is that advanced indexing always returns a copy of the data, not a view. This means changes to the new array created by advanced indexing will not affect the original array.

There are two main types of advanced indexing:

Integer Array Indexing
Boolean Array Indexing

Integer Array Indexing

This allows selection of arbitrary items in an array based on their N-dimensional coordinates, specified using integer arrays.

If you provide as many index arrays as there are dimensions in the target array, NumPy pairs corresponding elements from each index array to form the coordinates of the selected elements.

Example 1: Selecting specific elements

import numpy as np

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

# Define arrays for row and column indices
row_indices = np.array([0, 1, 2])
col_indices = np.array([0, 1, 0]) # Select elements at (0,0), (1,1), (2,0)

y = x[row_indices, col_indices]

print(f"Original array (x):\n{x}\n")
print(f"Selected elements (y): {y}")

# Modify the new array 'y'. This does NOT change 'x' because 'y' is a copy.
y[0] = 99
print(f"Modified selected elements (y): {y}")
print(f"Original array (x) remains unchanged:\n{x}")

Output:

Original array (x):
[[1 2 3]
 [4 5 6]
 [7 8 9]]

Selected elements (y): [1 5 7]
Modified selected elements (y): [99  5  7]
Original array (x) remains unchanged:
[[1 2 3]
 [4 5 6]
 [7 8 9]]

The shape of the resulting array matches the shape of the broadcasted index arrays.

Example 2: Selecting corners of a 2D array

import numpy as np

x = np.arange(12).reshape(4, 3)
print(f"Original array (x):\n{x}\n")

# Indices for top-left, top-right, bottom-left, bottom-right corners
# (0,0), (0,2), (3,0), (3,2)
row_indices = np.array([[0, 0], [3, 3]]) # Shape (2, 2)
col_indices = np.array([[0, 2], [0, 2]]) # Shape (2, 2)

corners = x[row_indices, col_indices]

print(f"Corner elements (shape matches index arrays):\n{corners}")

Output:

Original array (x):
[[ 0  1  2]
 [ 3  4  5]
 [ 6  7  8]
 [ 9 10 11]]

Corner elements (shape matches index arrays):
[[ 0  2]
 [ 9 11]]

Example 3: Combining Integer Indexing with Slicing

You can mix basic slicing with integer array indexing. Slicing selects ranges, while integer arrays select specific indices within those ranges (or other dimensions).

import numpy as np

x = np.arange(12).reshape(4, 3)
print(f"Original array (x):\n{x}\n")

# Select rows 1 and 3, and column 2
y = x[[1, 3], 2]
print(f"Rows 1 & 3, column 2: {y}")

# Select row 2, and columns 0 and 2
z = x[2, [0, 2]]
print(f"Row 2, columns 0 & 2: {z}")

# Select all rows (:), and columns 0 and 2
w = x[:, [0, 2]]
print(f"All rows, columns 0 & 2:\n{w}")

Output:

Original array (x):
[[ 0  1  2]
 [ 3  4  5]
 [ 6  7  8]
 [ 9 10 11]]

Rows 1 & 3, column 2: [ 5 11]
Row 2, columns 0 & 2: [6 8]
All rows, columns 0 & 2:
[[ 0  2]
 [ 3  5]
 [ 6  8]
 [ 9 11]]

Boolean Array Indexing

This involves using a boolean array (a mask) of the same shape as the original array. Elements corresponding to True values in the mask are selected, and those corresponding to False are discarded. The result is typically a 1-D array containing only the selected elements.

Example 1: Selecting elements based on a condition

import numpy as np

x = np.arange(12).reshape(4, 3)
print(f"Original array (x):\n{x}\n")

# Create a boolean mask for elements greater than 5
mask = x > 5
print(f"Boolean mask (x > 5):\n{mask}\n")

# Apply the mask
selected = x[mask]
print(f"Elements greater than 5: {selected}")

# You can apply the condition directly
print(f"Elements less than or equal to 5: {x[x <= 5]}")

Output:

Original array (x):
[[ 0  1  2]
 [ 3  4  5]
 [ 6  7  8]
 [ 9 10 11]]

Boolean mask (x > 5):
[[False False False]
 [False False False]
 [ True  True  True]
 [ True  True  True]]

Elements greater than 5: [ 6  7  8  9 10 11]
Elements less than or equal to 5: [0 1 2 3 4 5]

Example 2: Handling NaN values

Boolean indexing is useful for filtering out NaN (Not a Number) values.

import numpy as np

a = np.array([np.nan, 1., 2., np.nan, 3., 4., 5.])
print(f"Array with NaNs: {a}")

# Create mask for non-NaN values using np.isnan and the ~ (NOT) operator
mask = ~np.isnan(a)
print(f"Is not NaN mask: {mask}")

# Select non-NaN elements
filtered_a = a[mask]
print(f"Array without NaNs: {filtered_a}")

# Direct filtering
print(f"Directly filtered: {a[~np.isnan(a)]}")

Output:

Array with NaNs: [nan  1.  2. nan  3.  4.  5.]
Is not NaN mask: [False  True  True False  True  True  True]
Array without NaNs: [1. 2. 3. 4. 5.]
Directly filtered: [1. 2. 3. 4. 5.]

Example 3: Filtering based on data type

You can use functions like np.iscomplex to create masks for specific data types.

import numpy as np

a = np.array([1, 2 + 6j, 5, 3.5 + 5j, 0j])
print(f"Mixed type array: {a}")

# Select only complex numbers
complex_mask = np.iscomplex(a)
print(f"Is complex mask: {complex_mask}")
print(f"Complex elements only: {a[complex_mask]}")

Output:

Mixed type array: [1. +0.j 2. +6.j 5. +0.j 3.5+5.j 0. +0.j]
Is complex mask: [False  True False  True False]
Complex elements only: [2. +6.j 3.5+5.j]

Application: Advanced indexing is essential for data filtering, conditional selection and modification, rearranging elements, and implementing complex algorithms that require accessing specific subsets of array data.

For a comprehensive overview, consult the NumPy documentation on Indexing.