使用Python中NumPy库的where函数实现条件筛选

使用Python中NumPy库的where函数实现条件筛选

在Python的科学计算领域，NumPy是一款不可或缺的库。其中，where函数可以根据特定条件来进行筛选操作，非常方便实用。下面来通过几个示例演示where函数的使用方法。

首先，我们需要导入NumPy库：

import numpy as np

接下来，我们定义一个数组变量arr，用于演示where函数的使用：

arr = np.array([1, 2, 3, 4, 5])

示例一：索引筛选

使用where函数可以轻松实现基于索引的筛选，例如只选取arr中大于等于3的元素：

result = np.where(arr >= 3)
print(result)

输出结果：

(array([2, 3, 4], dtype=int64),)

这里的输出结果是一个tuple类型，第一个元素是满足条件的数组元素的索引。因此，我们可以通过下标访问到符合条件的元素值：

print(arr[result])

输出结果：

[3 4 5]

示例二：替换筛选

where函数还可以直接替换满足条件的元素值，例如将arr中大于等于3的元素都替换成0：

np.whe