Python numpy.argpartition函数方法的使用

NumPy 的 argpartition 函数用于将数组沿指定轴进行部分排序,并将元素的索引(partition)返回到一个新数组中。它不会改变原始数组,而是返回一个包含元素排序后索引的新数组。本文主要介绍一下NumPy中argpartition方法的使用。

numpy.argpartition

numpy.argpartition(a, kth, axis=-1, kind='introselect', order=None) [source]

使用kind关键字指定的算法,沿给定轴执行间接分区。 它沿着给定的轴按分区顺序返回与该索引数据具有相同形状的索引数组。

1.8.0版中的新功能。

参数:

a:array_like

数组进行排序。

kth: int 或 int的sequence

分区依据的元素索引。 

第k个元素将处于其最终排序位置,

所有较小的元素将在其之前移动,

而所有较大的元素将在其后面移动。 

分区中所有元素的顺序是不确定的。

 如果提供第k个序列,它将立即将它们全部划分到其排序位置。

axisintNone, 可选

要排序的轴。 默认值为-1(倒数第一个轴)。

 如果为None,则使用展平的数组。

kind{‘introselect’}, 可选

选择算法。 默认值为‘introselect’

orderstr 或 str的list, 可选

当a是定义了字段的数组时,

此参数指定要比较的字段的第一个,

第二个等。单个字段可以指定为字符串,

并且不需要指定所有字段,

但是仍将使用未指定的字段。 

他们以dtype出现的顺序打破关系。


返回值:        

index_arrayndarray, int

沿指定轴划分a的索引数组。 如果a是一维的,

a[index_array]产生分区a。 更一般而言,

np.take_along_axis(a,index_array,axis = 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)

推荐阅读
cjavapy编程之路首页