Функция argpartition

Функция argpartition выполняет частичную сортировку массива и возвращает индексы элементов, которые заняли бы указанные позиции при полной сортировке. Это похоже на partition, но возвращаются не сами элементы, а их индексы. Первым параметром передаётся массив, вторым - позиция (или несколько позиций), которые нужно частично отсортировать. Третьим параметром можно указать ось, по которой выполняется сортировка.

Синтаксис

np.argpartition(arr, kth, axis=-1, kind='introselect', order=None)

Пример

Давайте выполним частичную сортировку одномерного массива и получим индексы элементов для третьей позиции:

import numpy as np

arr = np.array([7, 2, 9, 1, 6, 3, 8, 5, 4])
res = np.argpartition(arr, 3)
print(res)

Результат выполнения кода:

[3 1 5 8 6 4 2 7 0]

В результате индексы первых четырёх элементов (до позиции 3) соответствуют наименьшим значениям массива, но порядок внутри этих элементов не гарантируется.

Пример

Проверим, что элементы с индексами из первых четырёх позиций действительно являются наименьшими в массиве:

import numpy as np

arr = np.array([7, 2, 9, 1, 6, 3, 8, 5, 4])
idx = np.argpartition(arr, 3)
print(arr[idx[:4]])

Результат выполнения кода:

[1 2 3 4]

Как видите, это действительно четыре наименьших элемента массива.

Пример

Можно указать несколько позиций для частичной сортировки. Например, получим индексы для двух позиций одновременно:

import numpy as np

arr = np.array([7, 2, 9, 1, 6, 3, 8, 5, 4])
res = np.argpartition(arr, [2, 5])
print(res)

Результат выполнения кода:

[3 1 5 8 6 4 2 7 0]

В этом случае элементы с индексами до позиции 2 будут наименьшими, а элементы с индексами до позиции 5 будут меньше всех остальных.

Пример

Давайте рассмотрим работу функции с двумерным массивом, указав ось для сортировки:

import numpy as np

arr = np.array([[7, 2, 9],
				[1, 6, 3],
				[8, 5, 4]])
res = np.argpartition(arr, 1, axis=0)
print(res)

Результат выполнения кода:

[[1 0 1]
 [0 2 2]
 [2 1 0]]

Здесь частичная сортировка выполняется по столбцам (ось 0). Каждый столбец сортируется независимо, и возвращаются индексы строк.

Пример

Используем функцию для сортировки по строкам двумерного массива:

import numpy as np

arr = np.array([[7, 2, 9],
				[1, 6, 3],
				[8, 5, 4]])
res = np.argpartition(arr, 0, axis=1)
print(res)

Результат выполнения кода:

[[1 0 2]
 [0 2 1]
 [2 1 0]]

Каждая строка сортируется отдельно, и возвращаются индексы столбцов для наименьшего элемента в каждой строке.

Пример

Рассмотрим использование параметра order для сортировки структурированных массивов:

import numpy as np

dtype = [('name', 'U10'), ('age', 'i4'), ('score', 'f4')]
arr = np.array([('Alice', 25, 85.5),
				('Bob', 30, 92.0),
				('Charlie', 22, 78.3),
				('David', 28, 88.7)],
				dtype=dtype)
res = np.argpartition(arr, 1, order='score')
print(res)

Результат выполнения кода:

[2 0 3 1]

Функция вернула индексы, где элемент с индексом 2 (Charlie) имеет наименьший score, а остальные индексы не гарантируют строгого порядка.

Пример

Применим функцию для поиска индексов наибольших элементов с использованием отрицательных значений:

import numpy as np

arr = np.array([7, 2, 9, 1, 6, 3, 8, 5, 4])
res = np.argpartition(-arr, 2)
print(res)

Результат выполнения кода:

[2 6 7 8 5 4 1 0 3]

Первые три индекса соответствуют элементам с наибольшими значениями (9, 8, 7), хотя порядок между ними не гарантируется.

Смотрите также

функцию partition,
которая возвращает частично отсортированные значения
функцию argsort,
которая возвращает индексы для полной сортировки
функцию sort,
которая выполняет полную сортировку массива
функцию searchsorted,
которая находит позиции для вставки элементов в отсортированный массив