Python numpy argpartition用法及代碼示例

本文簡要介紹 python 語言中 numpy.argpartition 的用法。

用法:

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

使用 kind 關鍵字指定的算法沿給定軸執行間接分區。它以分區順序沿給定軸返回與該索引數據具有相同形狀的索引數組。

參數：

a： array_like

要排序的數組。

kth： int 或整數序列

要分區的元素索引。 k-th 元素將處於其最終排序位置，所有較小的元素都將移到它之前，所有較大的元素都將移到它後麵。分區中所有元素的順序未定義。如果提供了 k-th 的序列，它將立即將所有這些分區到它們的排序位置。

axis： int 或無，可選

要排序的軸。默認值為 -1(最後一個軸)。如果為 None，則使用展平數組。

kind： {‘introselect’}，可選

選擇算法。默認為‘introselect’

order： str 或 str 列表，可選

當 a 是定義了字段的數組時，該參數指定首先比較哪些字段，第二個等。單個字段可以指定為字符串，不需要指定所有字段，但仍會使用未指定的字段，在他們在 dtype 中出現的順序，以打破關係。

返回：

index_array： ndarray，int

分區索引數組a沿著指定的軸。如果a是一維的，a[index_array]產生一個分區a.更普遍，np.take_along_axis(a, index_array, axis=a)總是產生分區的a，與維度無關。

注意：

有關不同選擇算法的說明，請參見 partition 。

例子：

一維數組：

>>> x = np.array([3, 4, 2, 1])
>>> x[np.argpartition(x, 3)]
array([2, 1, 3, 4])
>>> x[np.argpartition(x, (1, 3))]
array([1, 2, 3, 4])

>>> x = [3, 4, 2, 1]
>>> np.array(x)[np.argpartition(x, 3)]
array([2, 1, 3, 4])

多維數組：

>>> x = np.array([[3, 4, 2], [1, 3, 1]])
>>> index_array = np.argpartition(x, kth=1, axis=-1)
>>> np.take_along_axis(x, index_array, axis=-1)  # same as np.partition(x, kth=1)
array([[2, 3, 4],
       [1, 1, 3]])