Numpy 数组中获取符合条件的值

在使用numpy数组进行数据分析时，需要处理各种数据类型和数值，筛选出符合条件的数据是常见的操作之一。本文将介绍如何在numpy数组中获取符合条件的值。

阅读更多：Numpy 教程

1. numpy.where()函数

numpy.where函数是一种基本的查找函数，可以按条件在数组中查找出符合条件的数值，并将其输出。

1.1 简单用法

简单用法是在一个numpy数组中查询等于某个特定值的元素。

import numpy as np

a = np.array([1, 2, 3, 4])
filter = np.where(a == 2) # 查找等于2的元素
print(filter)

上述代码会输出如下结果：

(array([1]),)

1.2 复杂条件

numpy.where也可以用于查找复杂条件下的元素，比如在数组中查找大于某个特定值的元素。

import numpy as np

a = np.array([1, 2, 3, 4])
filter = np.where(a > 2) # 查找大于2的元素
print(filter)

上述代码会输出如下结果：

(array([2, 3]),)

1.3 处理多个条件

numpy.where也可以处理多个条件，比如在数组中查询大于2且小于5的元素。

import numpy as np

a = np.array([1, 2, 3, 4])
filter = np.where((a > 2) & (a < 5)) # 查找大于2且小于5的元素
print(filter)

上述代码会输出如下结果：

(array([2, 3]),)

2. 使用布尔运算符

numpy数组与布尔运算符结合使用，可以很方便地过滤出符合特定条件的数组元素。

import numpy as np

a = np.array([1, 2, 3, 4])
filter = a > 2 # 将大于2的元素筛选出来
print(a[filter])

上述代码会输出如下结果：

[3 4]

3. 筛选非NAN元素

numpy数组中存在一些空值NaN，如果需要排除这些NaN，可以使用numpy.isnan函数。

import numpy as np

a = np.array([1, np.nan, 3, 4])
filter = ~np.isnan(a) # 过滤NaN元素
print(a[filter])

上述代码会输出如下结果：

[1. 3. 4.]

4. numpy.compress()函数

numpy.compress函数用于根据特定条件选择序列中的元素，并返回新序列。函数的语法如下：

numpy.compress(condition, a, axis=None, out=None)

其中：

condition：bool类型的条件表达式，表示筛选条件。
a：待筛选的序列。
axis：沿着什么轴进行筛选。
out：输出结果。

4.1 沿着一维度进行筛选

import numpy as np

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

condition = [True, False, True, False, True, False] # 筛选条件
new_arr = np.compress(condition, arr) # 对数组进行筛选

print(new_arr)

上述代码会输出如下结果：

[1 3 5]

4.2 沿着多维度进行筛选

import numpy as np

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

condition = [True, False, True] # 筛选条件
new_arr = np.compress(condition, arr, axis=0) # 沿着列进行筛选

```python
print(new_arr)

上述代码会输出如下结果：

[[1 2 3]
 [7 8 9]]

5. numpy.extract()函数

numpy.extract函数与numpy.compress函数类似，不同之处在于它返回符合条件的元素，而不是在原始序列中的索引。

import numpy as np

arr = np.array([1, 2, 3, 4, 5, 6])
condition = [True, False, True, False, True, False]
new_arr = np.extract(condition, arr)

print(new_arr)

上述代码会输出如下结果：

[1 3 5]

6. numpy.ma库的使用

numpy.ma库专门处理有缺失值的数组数据，可以通过掩模数组(masked array)方式实现数组中的部分元素。

import numpy as np
import numpy.ma as ma

arr = np.array([1, 2, 3, 4, -999, 6])
mask = arr < 0   # 创建掩模数组
arr = ma.masked_array(arr, mask)

print(arr)
print(arr.mean())