在 NumPy 中,布尔索引 (boolean indexing) 是一种非常强大的数据选择方法。它可以让我们很方便地基于 “布尔数组” 来从另一个数组中选择元素。
什么是 NumPy 布尔索引?
布尔索引,指的是使用一个与原数组形状相同的布尔数组,来选择原数组中对应位置为元素。它最终会返回一个新的、扁平化的一维数组,其中包含了原数组中所有对应于布尔数组中 True 值的元素。
想要在 NumPy 中使用布尔索引,其实非常简单,只需要以下 2 步即可。
- 创建布尔数组:首先,我们需要创建一个与要索引的 NumPy 数组形状相同的布尔数组。这个布尔数组的每个元素都对应于原数组中的一个元素。
- 使用布尔数组:将这个布尔数组作为索引传递给原数组。NumPy 将会返回一个新的数组,其中只包含原数组中对应于布尔数组中 True 值的元素。
在使用 NumPy 布尔索引时,小伙伴们还要注意以下 3 点。
- 布尔数组与原数组形状必须完全相同,否则 NumPy 会报错。
- 通过布尔索引得到的是新数组,修改它不会影响原数组。
- 在使用逻辑运算符时,强烈建议加上括号。比如 arr > 10 & arr < 30 会因优先级错误而报错,正确写法为:(arr > 10) & (arr < 30)。
如何使用 NumPy 布尔索引?
接下来,我们通过几个简单的例子来讲解一下 NumPy 布尔索引是如何使用的。
示例 1:布尔索引的基本用法
import numpy as np
arr = np.array([10, 20, 30, 40, 50])
bool_arr = np.array([True, False, True, False, True])
result = arr[bool_arr]
print(result)运行结果如下。
[10 30 50]分析:
在这个例子中,布尔数组 bool_arr 的长度与 arr 的长度相同。然后布尔索引只会选择对应位置为 True 的元素。
示例 2:使用条件表达式
import numpy as np
arr = np.array([15, 22, 18, 27, 31, 20])
# 创建一个布尔数组,其中元素大于 20 的位置为 True
bool_arr = arr > 20
result = arr[bool_arr]
print(bool_arr)
print(result)运行结果如下。
[False True False True True False]
[22 27 31]分析:
在这个例子中,我们使用条件表达式 arr > 20 创建了一个布尔数组,此时得到 [False, True, False, True, True, False]。然后使用这个布尔数组来索引 arr,最终得到了所有大于 20 的元素。
示例 3:多个条件
import numpy as np
arr = np.array([10, 25, 15, 30, 35, 40])
# 找出大于 15 且小于 35 的元素
bool_arr = (arr > 15) & (arr < 35)
result = arr[bool_arr]
print(bool_arr)
print(result)运行结果如下。
[False True False True False False]
[25 30]分析:
条件表达式可以是一个条件,也可以是多个条件。在这个例子中,(arr > 15) & (arr < 35) 会创建一个布尔数组,即:[False True False True False False]。然后使用这个布尔数组来索引 arr,最终得到了所有大于 15 且小于 35 的元素。
示例 4:二维数组使用布尔索引
import numpy as np
arr = np.array([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
# 创建一个针对第一列的布尔数组
bool_arr = arr[:, 0] > 3
# 使用该布尔数组索引原始二维数组的行
result = arr[bool_arr]
print(result)运行结果如下。
[[4 5 6]
[7 8 9]]分析:
arr[:, 0] 会选中二维数组的 “第一列” :[1, 4, 7]。而 arr[:, 0] > 3 会生成布尔数组 [False, True, True],其中:
- 第 1 行的第一列值 1 不大于 3 → False。
- 第 2 行的第一列值 4 大于 3 → True。
- 第 3 行的第一列值 7 大于 3 → True。
arr[bool_arr] 会保留布尔数组为 True 的那些行,于是结果就是:
[[4 5 6]
[7 8 9]]