numpy.argpatition() 应用详解 + TopK值

numpy.argpartition(a, kth, axis=-1, kind=‘introselect’, order=None)
（一）一般numpy中以arg开头的函数都是返回下标，而不改变原数组。

• 参数a是指传入的Numpy.array
• 参数kth是指列表中下标为k位置应该放置该数组中第k大的值

例1：

import numpy as np
arr = np.array([3, 9, 1, 0, 2, 1, 7, 5])

# 返回的是下标，由于python是从0开始计算而非1，所以kth=0表示第0位置应当放置
# 该数组中最小的值
idx = np.argpartition(arr, 0)

[3 1 2 0 4 5 6 7]

# 经过排序的数组
arr[np.argpartition(arr, 0)]

[0 9 1 3 2 1 7 5]

（二）np.argpartition()将传入的数组arr分成两部分，即：排在第k位置前面的数都小于k，排在第k位置后面的值都大于k。

例2：

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

# 经过排序的数组
arr[np.argpartition(arr, 0)]

[0 9 1 3 2 1 7 5]

# 经过排序的数组，此时下标为1位置的值为该数组中第2小的值，排在下标1位置前的值
# 小于下标为1位置的值，排在下标1位置后的值大于下标为1位置的值
arr[np.argpartition(arr, 1)]

[0 1 9 3 2 1 7 5]

# 经过排序的数组
arr[np.argpartition(arr, 2)]

[0 1 1 3 2 9 7 5]

# 经过排序的数组
arr[np.argpartition(arr, 3)]

[0 1 1 2 3 9 7 5]

# 经过排序的数组
arr[np.argpartition(arr, 4)]

[2 0 1 1 3 9 7 5]

（三）TopK值

例3：获得最大的4个值（k=4）

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

k = 4 

# 将该数组中第4大的值放置在倒数第4的位置
idx = np.argpartition(arr, -k)
new_arr = arr[idx]
max_4 = new_arr[-k:]