Python中的np.where()和nan

介绍

在Python的NumPy库中，np.where()函数可以让我们找到满足条件的元素的索引或替换数组中的元素。同时，当遇到NaN（Not a Number）时，该函数有不同的行为。本文将介绍np.where()在处理数组中的NaN时的情况。

np.where()基本用法

np.where()有以下两个基本用法：

1. 根据条件返回符合条件的元素的索引。
2. 根据条件替换数组中符合条件的元素。

1. 根据条件返回符合条件的元素的索引

import numpy as np

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

print(idx)

输出：

(array([3, 4, 5]),)

2. 根据条件替换数组中符合条件的元素

import numpy as np

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

print(new_arr)

输出：

[0 0 0 4 5 6]

处理NaN

当数组中存在NaN时，np.where()有不同的行为。下面将分别介绍。

1. 根据条件返回符合条件的元素的索引

当数组中存在NaN时，np.where()仍然可以找到满足条件的元素的索引，但是对于NaN，它返回的是一个空元组。

import numpy as np

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

print(idx)

输出：

(array([3, 4, 5]),)

但是如果我们想要找到数组中的NaN元素，我们可以使用np.isnan()函数。

import numpy as np

arr = np.array([1, 2, np.nan, 4, 5, 6])
idx = np.where(np.isnan(arr))

print(idx)

输出：

(array([2]),)

2. 根据条件替换数组中符合条件的元素

当数组中存在NaN时，np.where()在替换元素时，会将NaN替换成指定的值。

import numpy as np

arr = np.array([1, 2, np.nan, 4, 5, 6])
new_arr = np.where(arr < 4, 0, arr)

print(new_arr)

输出：

[1. 2. nan 4. 5. 6.]

可以看到，由于NaN无法与任何数值进行比较，因此在替换时保持不变。

如果我们想要将NaN替换为指定的值，可以先使用np.isnan()函数找到NaN的索引，然后再将这些索引对应的元素替换成新值。

import numpy as np

arr = np.array([1, 2, np.nan, 4, 5, 6])
nan_idx = np.where(np.isnan(arr))
arr[nan_idx] = 0

print(arr)

输出：

[1. 2. 0. 4. 5. 6.]

可以看到，我们成功将数组中的NaN替换成了0。

结论

在处理数组中的NaN时，我们应该谨慎地使用np.where()函数。在找到满足条件的元素的索引时，我们需要使用np.isnan()函数，而在替换数组中的元素时，我们需要先找到NaN元素的索引，然后再进行替换。