numpy数组最常用的4个搜索方法

NumPy提供了一些搜索和查找数组中元素的方法，包括：

1. np.where(condition[, x, y])：返回满足条件的元素的下标。可以指定x和y参数，如果不指定，则返回元素下标。
2. np.argwhere(condition)：返回满足条件的元素的下标，与where()方法类似，但返回的是一个包含下标的数组，而不是元组。
3. np.searchsorted(arr, v[, side, sorter])：返回将值v插入到有序数组arr中的位置，可以指定搜索方向和排序方法。
4. np.extract(condition, arr)：返回满足条件的元素，以一维数组的形式返回。

下面分别进行详细介绍：

np.where(condition[, x, y])

np.where()方法的第一个参数是条件，可以是一个数组或一个条件表达式。它会返回一个由满足条件的元素下标构成的元组或数组，如果指定了x和y参数，则满足条件的元素会被替换为x中的元素，否则将被替换为元素下标。

import numpy as np

# 生成一个随机数组
a = np.random.randint(1, 10, size=(3, 4))
print(a)
# 找到大于5的元素的下标
index = np.where(a > 5)
print(index)
# 替换大于5的元素为0
a[np.where(a > 5)] = 0
print(a)

输出结果：

[[2 2 4 7]
 [4 8 7 1]
 [1 9 9 9]]
(array([0, 1, 1, 2, 2, 2, 2]), array([3, 1, 2, 1, 2, 3, 4]))
[[2 2 4 0]
 [4 0 0 1]
 [1 0 0 0]]

np.argwhere(condition)

np.argwhere()方法返回满足条件的元素的下标，与np.where()方法类似，但返回的是一个包含下标的数组，而不是元组。

示例：

import numpy as np

# 生成一个随机数组
a = np.random.randint(1, 10, size=(3, 4))
print(a)
# 找到大于5的元素的下标
index = np.argwhere(a > 5)
print(index)

输出结果：

[[6 1 6 8]
 [5 5 1 7]
 [2 8 7 5]]
[[0 3]
 [1 0]
 [1 1]
 [1 3]
 [2 1]
 [2 2]]

np.searchsorted(arr, v[, side, sorter])

searchsorted()函数的返回值为一个整数或整数数组，表示插入元素后应该在哪个位置。

如果v的值小于a中的最小元素，则返回0；如果v的值大于a中的最大元素，则返回a的长度。

例如：

import numpy as np

a = np.array([1, 2, 3, 4, 5])
v = 3.5

print(np.searchsorted(a, v, side='left'))  # 输出：3
print(np.searchsorted(a, v, side='right'))  # 输出：3

extract(condition,arr)函数

extract()函数根据指定的条件，从数组中提取满足条件的元素。它返回一个满足条件的一维数组。

下面是一个使用extract()函数的示例：

import numpy as np

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

print(np.extract(condition, arr))  # 输出：[3 4 5]