numpy中的argpartition用法

一、np.argpartion()
Numpy的argpartion函数来源于快排算法中的一个典型操作partition，即根据一个数值x，把数组中的元素划分成两半，使得index前面的元素都不大于x，index后面的元素都不小于x。
np.argpartition不对原数组操作，只返回一个按照上述操作排序过后的index。通过这个函数可以高效地找到 N 个最大值的索引并返回 N 个值。在给出索引后，我们可以根据需要进行值排序。
 

import numpy as np
array = np.array([10, 7, 4, 3, 2, 2, 5, 9, 0, 4, 6, 0])
#返回一个索引，比原数组第5大（从0开始）的数小的数在这个数之前，比这个数大的数在它之后。
index = np.argpartition(array, 4)
#输出，新索引
print(index)
#[ 4 11  8  5  3  2  9  6  1 10  7  0]
#按这个新索引可以重新排列数组
print(array[index])
#[ 2  0  0  2  3  4  4  5  7  6  9 10]
#第5大的数是3，比3小的在3之前，比3大的在3之后

二、输出top5

np.argpartition的一个重要应用就是高效输出最大的几个值(如top-5)，因为不用像np.sort对所有元素排序

#还是上边那个数组，输出top5
array = np.array([10, 7, 4, 3, 2, 2, 5, 9, 0, 4, 6, 0])
array[np.argpartition(array, -5)[-5:]]
#输出：[ 5,  7,  6,  9, 10]

import numpy as np

arr = np.array([66, 15, 27, 33, 19, 13, 10])

"""
np.partition(arr, n)
找出arr中第n + 1小的元素(将arr排序之后索引n的元素), 然后返回一个新数组
并将原来数组中第n + 1小的元素放在新数组索引为n的地方, 保证左边的元素比它小, 右边的元素比它大
"""
print(np.partition(arr, 3))  # [15 13 10 19 27 33 66]
# 第4小的元素(排完序之后索引为3)显然是19, 那么将19放在索引为3的位置, 然后左边的元素比它小, 右边的元素比它大
# 至于两边的顺序则没有要求

# 虽然我们可以使用sort, 但是sort是全局排序
# 如果数组非常大, 我们只希望选择最小的10个元素, 直接通过np.partition(arr, 9)即可
# 然后如果排序的话, 只对这选出来的10个元素排序即可, 而无需对整个大数组进行排序

# 同理还可以从后往前找, 比如:
# np.partition(arr, -2)表示找到第2大的元素(将arr排序之后索引-2的元素), 放在数组索引为-2的地方
# 然后左边元素比它小, 右边元素比它大
print(np.partition(arr, -2))  # [13 10 27 15 19 33 66]
# 第2大的元素显然是33, 那么排在索引为-2的位置, 左边元素比它小, 右边元素比它大


# 然后argpartition不用想, 肯定是获取排序之后的索引
print(np.argpartition(arr, 3))  # [1 5 6 4 2 3 0]
print(np.argpartition(arr, -2))  # [5 6 2 1 4 3 0]

例子一：找到数组的最大最小值：

>>> import numpy as np
>>> arr = np.array([46, 57, 23, 39, 1, 10, 0, 120])
>>> np.argpartition(arr, len(arr)-1)
array([5, 3, 2, 6, 4, 0, 1, 7], dtype=int32)
>>> arr[np.argpartition(arr, len(arr)-1)[len(arr)-1]]
120
>>> arr[np.argpartition(arr, 0)[0]]
0

例子二：找到数组的第3小（index=2）的值和第2大（index=-2）的值。

>>> arr[np.argpartition(arr, 2)[2]]
10
>>> arr[np.argpartition(arr, -2)[-2]]
57

例子三：同时找到第3和第4小的值。注意这里，用[2,3]同时将第3和第4小的排序好，然后可以分别通过下标[2]和[3]取得。

>>> arr[np.argpartition(arr, [2,3])[2]]
10
>>> arr[np.argpartition(arr, [2,3])[3]]
23