numpy.argpartition
numpy.argpartition(a, kth, axis=-1, kind='introselect', order=None) [source]
使用kind关键字指定的算法，沿给定轴执行间接分区。 它沿着给定的轴按分区顺序返回与该索引数据具有相同形状的索引数组。
1.8.0版中的新功能。
参数: | a：array_like 数组进行排序。 kth： 分区依据的元素索引。 第k个元素将处于其最终排序位置， 所有较小的元素将在其之前移动， 而所有较大的元素将在其后面移动。 分区中所有元素的顺序是不确定的。 如果提供第k个序列，它将立即将它们全部划分到其排序位置。 axis： 要排序的轴。 默认值为 如果为 kind： 选择算法。 默认值为 order： 当a是定义了字段的数组时， 此参数指定要比较的字段的第一个， 第二个等。单个字段可以指定为字符串， 并且不需要指定所有字段， 但是仍将使用未指定的字段。 他们以 |
返回值: | index_array: 沿指定轴划分a的索引数组。 如果a是一维的， 则
总是产生分区 |
Notes
有关不同选择算法的说明，请参见partition。
例子，
1）一维数组
import numpy as np # 创建一个NumPy数组 array = np.array([5, 2, 4, 1, 3]) # 获取前3个元素的下标 k = 3 indices = np.argpartition(array, k) # 打印排序后的索引 print(indices) # 获取前3个元素 top_3_elements = array[indices[:k]] # 打印前3个元素 print(top_3_elements)
2）多维数组
import numpy as np # 创建一个 NumPy 数组 x = np.array([[3, 4, 2], [1, 3, 1]]) # 使用 argpartition 函数获取每个子数组的第 2 大元素的索引 kth = 1 # 指定要获取的第 K 大元素的索引，这里是第 2 大元素 index_array = np.argpartition(x, kth=kth, axis=-1) # 使用 np.take_along_axis 函数根据索引重新排列元素，实现部分排序 sorted_array = np.take_along_axis(x, index_array, axis=-1) # 打印排序后的数组 print(sorted_array)