Python numpy.argpartition函数方法的使用

numpy.argpartition 是 NumPy 中用于部分排序索引的函数,主要用于高效地找出第 k 小的元素及其位置,而不进行完整排序。numpy.argpartition 非常适合用于从数组中快速找出前 k 小或前 k 大的元素。与 argsort 相比,它在处理大型数组时具有更高的效率,因为它只进行部分排序而非全排序。需要注意的是,argpartition 返回的前 k 个索引对应的是目标范围内的最小或最大元素,但这些元素在结果中的顺序并不保证是有序的。本文主要介绍一下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”

或der :str 或 str的类型list, 可选

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

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

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

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

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

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

返回值 :

index_array :ndarray, int

沿指定轴对数组 a 进行索引划分。

如果 a 是一维数组,

那么通过 a[index_array] 

就可以得到已分区的结果。

更一般地说,无论 a 是几维数组,

使用 np.take_along_axis(a, 

index_array, axis=...) 

都可以按给定的索引顺序

提取出分区后的结果。

例子

1)一维数组

import numpy as np

# 找出前 4 小的元素(即全部元素,但顺序不完全排序)
x = np.array([3, 4, 2, 1])
result1 = x[np.argpartition(x, 3)]
print(result1)
# 输出示例:array([2, 1, 3, 4])
# 注意:前 4 小的元素都包含,但它们的顺序未必有序

# 找出第 1 小到第 3 小的元素范围
result2 = x[np.argpartition(x, (1, 3))]
print(result2)
# 输出示例:array([1, 2, 3, 4])
# 说明:包含第 1 和第 3 小的元素,也可能正好是完全排序的结果

# 从列表类型中使用 argpartition(需要显式转换为数组)
x_list = [3, 4, 2, 1]
result3 = np.array(x_list)[np.argpartition(x_list, 3)]
print(result3)
# 输出示例:array([2, 1, 3, 4])

2)多维数组

import numpy as np

# 创建一个二维数组
x = np.array([[3, 4, 2], 
              [1, 3, 1]])

# 对每一行进行部分排序,找到每行中第1小的元素,并将其放在正确位置
index_array = np.argpartition(x, kth=1, axis=-1)

# 根据索引获取部分排序后的数组(同 np.partition(x, kth=1) 效果)
result = np.take_along_axis(x, index_array, axis=-1)

print(result)
# 输出:
# [[2 3 4]
#  [1 1 3]]

3)列表排序

import numpy as np

# 原始数组
a = np.array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])

# 使用 argpartition 找出前 5 小的元素位置(k = 4)
indices = np.argpartition(a, 4)
print("部分排序后的索引结果:", indices)

# 使用索引查看部分排序结果(前 5 小的元素在前面,但顺序不保证)
print("部分排序后的数组:", a[indices])

4)取列表中前5个元素 

import numpy as np

# 原始数组
a = np.array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])

# 找出最大的 5 个元素(顺序不保证)
top5 = a[np.argpartition(a, -5)[-5:]]

print(top5)

推荐阅读
cjavapy编程之路首页