Numpy 的一些以 arg 开头的函数
古明地觉约 6 分钟楔子
Numpy 里面有一些以 arg 开头的函数，非常的有意思，因为它们返回的不是元素、而是元素的索引，下面来看一下用法。
np.argmax
我们知道 np.max 是获取最大的元素，那么np.argmax是做什么的呢？
import numpy as nparr = np.array([3, 22, 4, 11, 2, 44, 9])
print(np.max(arr)) # 44
print(np.argmax(arr)) # 5
结论很清晰，np.argmax 是获取最大元素所在的索引。
同理还有 np.argmin，np.min 是获取数组中最小的元素，显然是2；np.argmin 是获取数组中最小元素对应的索引，显然是 4。
import numpy as nparr = np.array([3, 22, 4, 11, 2, 44, 9])
print(np.min(arr)) # 2
print(np.argmin(arr)) # 4
那么问题来了，如果我们想自己实现一个 argmax 和 argmin 该怎么做呢？
def argmax(lst):
lst = [
# 将每个元素和对应的索引组合起来
(item, idx) for idx, item in enumerate(lst)
]
return max(lst)[1]print(
argmax([3, 22, 4, 11, 2, 44, 9])
) # 5
还是很简单的，如果实现 argmin 的话，只需将里面的 max 函数换成 min 即可。
np.argwhere
np.where 还是用的挺频繁的，先来复习一下它的用法吧。
import numpy as nparr = np.array(
[1, 2, 3, 4, 5, 6, 7, 8])
# 如果元素大于 4，那么减去 10，否则乘以 10
print(
np.where(arr > 4, arr - 10, arr * 10)
) # [10 20 30 40 -5 -4 -3 -2]
# 如果元素大于4, 那么保持不变, 否则变成 4
print(
np.where(arr > 4, arr, 4)
) # [4 4 4 4 5 6 7 8]
和 np.where 作用类似的还有一个 np.clip，来看一下。
import numpy as nparr = np.array([1, 2, 3, 4, 5, 6, 7, 8])
# 小于 2 的换成 2, 大于 6 的换成 6
# 一般在设置上下限的时候非常有用
print(
np.clip(arr, 2, 6)
) # [2 2 3 4 5 6 6 6]
那么 np.argwhere 是做啥的呢？首先这个函数只接受一个参数，找出满足条件的元素对应的索引。
import numpy as nparr = np.array([3, 4, 5, 6, 7])
# 显然 3、5、7 对 2 取模不等于 0
# 它们对应的索引分别是 0、2、4
print(np.argwhere(arr % 2 != 0))
"""
[[0]
[2]
[4]]
"""
# 需要扁平化
print(
np.argwhere(arr % 2 != 0).flatten()
) # [0 2 4]
显然元素 3、5、7 在 %2 之后不为 0，所以会筛选出它们的索引，因此是 [0 2 4]。只不过默认不是一个一维数组，我们需要再调用一下 flatten，将其扁平化。
np.argsort
np.sort 是用来排序的，类似于内置函数 sorted。
import numpy as nparr = np.array([44, 22, 33, 66, 55, 11])
print(
np.sort(arr)
) # [11 22 33 44 55 66]
sort 很容易，再来看看 argsort。
import numpy as nparr = np.array([44, 22,33, 66, 55, 11])
print(
np.sort(arr)
) # [11 22 33 44 55 66]
print(
np.argsort(arr)
) # [5 1 2 0 4 3]
所以 argsort 返回的是元素排序后对应的索引，光说可能不好理解，我们画一张图就清晰了。
因此 sort 是对元素进行排序，而 argsort 是按照元素的大小对索引进行排序。比如我们想选择数组中最大的 3 个元素，就有两种做法。
import numpy as nparr = np.array([44, 22, 33, 66, 55, 11])
# 选择前三个最大元素，可以先全局排序
# 然后选择前三个
print(
np.sort(arr)[-3:]
) # [44 55 66]
# 使用 argsort，找到最大的三个元素的索引
# 然后通过 arr 进行筛选
print(
arr[np.argsort(arr)[-3:]]
) # [44 55 66]
另外 argsort 还可以实现一个特殊的需求，那就是排名统计，举个例子：[44, 22, 33, 66, 55, 11]，如果按照从小到大排名，
44 在 array 中排名第 4；
22 在 array 中排名第 2；
33 在 array 中排名第 3；
66 在 array 中排名第 6；
55 在 array 中排名第 5；
11 在 array 中排名第 1;
所以需要实现一个功能，将 arr 传进去之后，返回 [4 2 3 6 5 1]。如果是你的话，你会怎么做呢？最容易想到的办法就是遍历：
lst = [44, 22, 33, 66, 55, 11]
# 先排序
sorted_lst = sorted(lst)
# 依次查看 lst 里面的每个元素，在排序后列表中的位置
# 但索引从 0 开始，因此需要再加上 1
rank = [sorted_lst.index(item) + 1
for item in lst]
print(rank) # [4, 2, 3, 6, 5, 1]但是这段代码有缺陷，那就是 lst 里面的元素不能重复，这里我们假设不重复。如果用 argsort 该怎么做呢？
import numpy as nparr = np.array([44, 22, 33, 66, 55, 11])
print(
np.argsort(arr)
) # [5 1 2 0 4 3]
print(
np.argsort(np.argsort(arr))
) # [3 1 2 5 4 0]
print(
np.argsort(np.argsort(arr)) + 1
) # [4 2 3 6 5 1]
# np.argsort(arr) 等价于 arr.argsort()
使用 Numpy 的话两次 argsort 即可完成，这个过程可能有点绕，为了方便理解，我们画一张图：
可以对照着图多理解一下。
np.argpartition
介绍 argpartiiton 之前，先来说一下 partition，它和 np.sort 类似，都是用于排序，但 partiton 是局部排序。
import numpy as nparr = np.array([66, 15, 27, 33, 19, 13, 10])
"""
np.partition(arr, k) 返回一个新数组
并且索引为 k 的元素比它左边的元素都大，比它右边的元素都小
"""
print(
np.partition(arr, 3)
) # [15 13 10 19 27 33 66]
# 返回的数组中索引为 3 的元素是 19
# 比它左边的元素都大，比它右边的元素都小
# 至于左右两边的顺序则没有要求
当我们想实现 topK 的时候，这个方法非常适合。
虽然我们可以使用 np.sort 排序之后再通过切片截取, 但 sort 是全局排序。如果数组非常大, 我们只希望选择最小的 10 个元素，那么全局排序就有点浪费了。而最好的方式是通过 np.partition(arr, 9)，然后将前 10 个元素截取出来进行排序即可，而无需对整个大数组进行排序。
结果是一样的，但是两者的效率如何呢？我们来对比一下：
我们看到效率差的不是一点点，原因就在于 partition 是一个时间复杂度为 O(N) 算法。因此在处理 topK 问题时，如果 K 远小于长度 N，那么使用 partition 是最佳选择。
另外，partition 的第二个参数还可以为负数：
import numpy as nparr = np.array([66, 15, 27, 33, 19, 13, 10])
# 让数组中索引为 -2 的元素比它左边的元素都大，
# 比它右边的元素都小，显然这是用来筛选前 K 个最大的元素
print(
np.partition(arr, -2)
) # [15 13 10 19 27 33 66]
# 比如筛选前 3 个最大的元素
print(
np.partition(arr, -3)[-3:]
) # [27 33 66]
partition 函数的用法我们就知道了，而 argpartition 也和 argsort 类似，只不过返回的是对应的索引。可以自己编写代码测试一下。
以上就是 arg 开头的几个函数的用法。