相比于常规的 Python 序列，NumPy 提供了更多了索引功能。除了前面看到的通过整数和切片索引，还可以通过整数数组和布尔数组对数组进行索引。
使用索引数组进行索引#
import numpy as np
a = np.arange(12)**2 # 最前面的 12 个平方数
i = np.array([1, 1, 3, 8, 5]) # 一个索引数组
print(a[i]) # [ 1 1 9 64 25]，数组 a 在位置 i 处的元素
j = np.array([[3, 4], [9, 7]]) # 二维的索引数组
print(a[j]) # 结果数组的形状与 j 相同
# [[ 9 16]
# [81 49]]
当被索引数组 a 为多维时，单个索引数组将引用 a 的第一维。以下示例通过使用调色板将将一个标签图像转换为一个实际的颜色图像。
import numpy as np
palette = np.array([[0, 0, 0], # 黑色
[255, 0, 0], # 红色
[0, 255, 0], # 绿色
[0, 0, 255], # 蓝色
[255, 255, 255]]) # 白色
image = np.array([[0, 1, 2, 0], # 每个值对应调色板中的一个颜色
[0, 3, 4, 0]]) # 一个包含 2*4 个像素的图像
print(palette[image]) # (2, 4, 3) 颜色图像
# [[[ 0 0 0]
# [255 0 0]
# [ 0 255 0]
# [ 0 0 0]]
# [[ 0 0 0]
# [ 0 0 255]
# [255 255 255]
# [ 0 0 0]]]
我们还可以给定超过一维的索引。每一维度的索引数组应具有相同的形状。
import numpy as np
a = np.arange(12).reshape(3, 4)
print(a)
# [[ 0 1 2 3]
# [ 4 5 6 7]
# [ 8 9 10 11]]
i = np.array([[0, 1], # 数组 a 的第一维索引
[1, 2]])
j = np.array([[2, 1], # 第二维索引
[3, 3]])
print(a[i, j]) # i 和 j 必须具有相同的形状
# [[ 2 5]
# [ 7 11]]
print(a[i, 2])
# [[ 2 6]
# [ 6 10]]
print(a[:, j])
# [[[ 2 1]
# [ 3 3]]
# [[ 6 5]
# [ 7 7]]
# [[10 9]
# [11 11]]]
在 Python 中，arr[i, j] 和 arr[(i, j)] 几乎是完全等同的，因此我们可以将 i 和 j 放入一个元组中，并使用如下方法进行索引。
import numpy as np
a = np.arange(12).reshape(3, 4)
print(a)
# [[ 0 1 2 3]
# [ 4 5 6 7]
# [ 8 9 10 11]]
i = np.array([[0, 1], # 数组 a 的第一维索引
[1, 2]])
j = np.array([[2, 1], # 第二维索引
[3, 3]])
l = (i, j) # 等同于 a[i, j]
print(a[l])
# [[ 2 5]
# [ 7 11]]
不过，我们并不能把 i 和 j 放进一个数组中，因为这个数组将被解释维对 a 的第一维数据的索引。
import numpy as np
a = np.arange(12).reshape(3, 4)
print(a)
# [[ 0 1 2 3]
# [ 4 5 6 7]
# [ 8 9 10 11]]
i = np.array([[0, 1], # 数组 a 的第一维索引
[1, 2]])
j = np.array([[2, 1], # 第二维索引
[3, 3]])
s = np.array([i, j])
# 不是我们想要的
# print(a[s]) # IndexError: index 3 is out of bounds for axis 0 with size 3
print(a[tuple(s)]) # 与 a[i, j] 相同
# [[ 2 5]
# [ 7 11]]
另一个使用数组进行索引的常见用法是搜索时间相关数据序列的最大值：
import numpy as np
time = np.linspace(20, 145, 5) # 时间标度
data = np.sin(np.arange(20)).reshape(5, 4) # 4 组时间相关数据
print(time) # [ 20. 51.25 82.5 113.75 145. ]
print(data)
# [[ 0. 0.84147098 0.90929743 0.14112001]
# [-0.7568025 -0.95892427 -0.2794155 0.6569866 ]
# [ 0.98935825 0.41211849 -0.54402111 -0.99999021]
# [-0.53657292 0.42016704 0.99060736 0.65028784]
# [-0.28790332 -0.96139749 -0.75098725 0.14987721]]
ind = data.argmax(axis=0) # 每组数据的最大值的索引
print(ind) # [2 0 3 1]
time_max = time[ind] # 最大值对应的时间
data_max = data[ind, range(data.shape[1])] # => data[ind[0], 0], data[ind[1], 1]...
print(time_max) # [ 82.5 20. 113.75 51.25]
print(data_max) # [0.98935825 0.84147098 0.99060736 0.6569866 ]
print(np.all(data_max == data.max(axis=0))) # True