NumPy 高级索引学习笔记

NumPy 高级索引指的是使用整数数组和布尔数组来访问数组中的元素。它与基本索引不同，可以用于获取、修改和生成新数组。

整数数组索引

使用整数数组作为索引值时，会获取对应索引位置上的元素并生成新数组。

pythonCopy Code
import numpy as np

# 创建数组
arr = np.array([1, 2, 3, 4, 5])

# 使用整数数组进行索引
idx = np.array([0, 2, 4])
new_arr = arr[idx]

print(new_arr)  # [1 3 5]

如果需要使用多维整数数组索引，需要在每个维度上分别传入一个整数数组。

pythonCopy Code
# 创建二维数组
arr = np.array([[1, 2], [3, 4], [5, 6]])

# 使用多维整数数组进行索引
idx = np.array([[0, 1], [1, 2]])
new_arr = arr[idx]

print(new_arr)  # [[1 2], [4 5]]

布尔数组索引

使用布尔数组作为索引值时，会获取等长的布尔数组中为 True 的元素所在位置上的元素，并生成新数组。

pythonCopy Code
# 创建数组
arr = np.array([1, 2, 3, 4, 5])

# 使用布尔数组进行索引
condition = np.array([False, True, False, True, False])
new_arr = arr[condition]

print(new_arr)  # [2 4]

如果需要在多维数组中使用布尔数组索引，可以分别传入每个维度上对应的布尔数组。

pythonCopy Code
# 创建二维数组
arr = np.array([[1, 2], [3, 4], [5, 6]])

# 使用布尔数组进行索引
condition = np.array([[False, True], [True, False], [True, True]])
new_arr = arr[condition]

print(new_arr)  # [2 3 5 6]

实例

假设我们有一个 $5\times5$ 的矩阵：

pythonCopy Code
matrix = np.arange(1, 26).reshape(5, 5)
print(matrix)

输出：

Copy Code
[[ 1  2  3  4  5]
 [ 6  7  8  9 10]
 [11 12 13 14 15]
 [16 17 18 19 20]
 [21 22 23 24 25]]

现在我们想要获取第一行、第三行和第四行的元素，可以使用整数数组索引：

pythonCopy Code
idx = np.array([0, 2, 3])
new_matrix = matrix[idx]
print(new_matrix)

输出：

Copy Code
[[ 1  2  3  4  5]
 [11 12 13 14 15]
 [16 17 18 19 20]]

又假设我们想要获取矩阵中大于 10 的元素，可以使用布尔数组索引：

pythonCopy Code
condition = matrix > 10
new_matrix = matrix[condition]
print(new_matrix)

输出：

Copy Code
[11 12 13 14 15 16 17 18 19 20 21 22 23 24 25]