numpy.argpartition

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

Функция argpartition() возвращает массив индексов элементов исходного массива в их разбиении по указанному значению.

Работа данной функции абсолютно аналогична numpy.partition(), но вместо копии разделенного исходного массива она возвращает массив индексов его элементов в этом разбиении.

Если исходный массив a является плоским, то a[np.argpartition(a)] вернет искомое разбиение. В случаях, когда a является многомерным, то получить искомое разбиение можно либо на основе итерирования, либо с помощью команды np.take_along_axis(a, np.argpartition(a, axis = ax), axis = ax).

Параметры:
a - массив NumPy или подобный массиву объект.
Исходный массив.
kth - целое число или последовательность целых чисел.
Индекс или последовательность индексов элементов в его отсортированной копии.
axis - целое число (необязательный параметр).
Определяет ось вдоль которой выполняется разбиение элементов. Если равен None, то разбиение выполняется по сжатому до одной оси представлению исходного массива. По умолчанию axis = -1, что соответствует разбиению по последней оси массива.
kind - строка 'introselect' (необязательный параметр).
Алгоритм выборки элементов. Единственное возможное значение 'introselect'.
order - строка или список строк (необязательный параметр).
В случае, когда массив a является структурированным, и его полям присвоено имя, то данный параметр позволяет указать порядок в котором они должны учавствовать в сортировке Если указаны не все поля, то неуказанные будут использоваться в том порядке в котором они перечислены в dtype структурированного массива.
Возвращает:
ndarray - массив NumPy
массив индексов, разделяющих исходный массив a, на подгруппы элементов больших и меньших указанных элементов.
Смотрите так же:
partition, sort

Замечание

Данная функция обладает эквивалентным методом класса ndarray, т.е. np.argpartition(a) равносильно вызову метода a.argpartition():

>>> import numpy as np
>>> 
>>> a = np.random.randint(0, 20, 10)
>>> a
array([19, 10, 16, 15, 19,  1, 11, 18,  3,  1])
>>> 
>>> np.argpartition(a, 5)
array([5, 8, 6, 9, 1, 3, 2, 7, 0, 4], dtype=int32)
>>> 
>>> a.argpartition(5)
array([5, 8, 6, 9, 1, 3, 2, 7, 0, 4], dtype=int32)

Используемый алгоритм выборки делает временную копию данных при выполнении сортировки по всем осям кроме последней, так что сортировка только по последней оси выполняется быстрее всего и занимает меньше всего места.

Сортировка в массивах содержащих действительные числа и nan всегда сначала выполняется над действительными числами, а уже затем над nan.

Сортировка комплексных чисел выполняется лексикографически: если действительная и мнимая часть не равна nan то сортировка выполняется по действительной части, но если действительные части равны, то порядок определяется их мнимыми частями. Расширенный порядок сортировки: [R + Rj, R + nan, nan + Rj, nan + nan], где R - действительное число.



Примеры

Для удобства мы создадим объект генератора случайных чисел и увеличим область вывода массивов:

>>> import numpy as np
>>> 
>>> rng = np.random.RandomState(777)
>>> np.set_printoptions(edgeitems = 10, linewidth = 200)

Теперь создадим массив случайных целых чисел:

>>> a = rng.randint(0, 1000, 12)
>>> a
array([103, 815, 571, 934, 599, 945,  71, 669, 639, 116, 423, 985])

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

>>> ind = np.argpartition(a, 5)
>>> ind
array([ 0,  9,  6, 10,  2,  4,  1,  7,  8,  3,  5, 11], dtype=int32)
>>> 
>>> a[ind]
array([103, 116,  71, 423, 571, 599, 815, 669, 639, 934, 945, 985])

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

>>> b = np.arange(12) % 2
>>> b
array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1], dtype=int32)
>>> 
>>> b[ind]
array([0, 1, 0, 0, 0, 0, 1, 1, 0, 1, 1, 1], dtype=int32)

Для многомерных массивов, получить разбиение по возвращаемому индексу проще всего с помощью функции np.take_along_axis():

>>> a = rng.randint(0, 100, size = (6, 6))
>>> a
array([[43,  0, 62, 38, 53, 78],
       [28, 92, 93, 97, 20, 23],
       [10, 74,  4, 80, 16, 85],
       [88, 15,  0, 62, 53, 55],
       [89, 54, 46, 61, 17,  8],
       [ 9, 16, 85, 62, 95, 54]])
>>> 
>>> np.partition(a, 3, axis = 0)
array([[ 9, 15,  0, 61, 17, 23],
       [10,  0,  4, 38, 16,  8],
       [28, 16, 46, 62, 20, 54],
       [43, 54, 62, 62, 53, 55],
       [89, 74, 93, 80, 53, 78],
       [88, 92, 85, 97, 95, 85]])
>>> 
>>> 
>>> ind = np.argpartition(a, 3, axis = 0)
>>> ind
array([[5, 3, 3, 4, 4, 1],
       [2, 0, 2, 0, 2, 4],
       [1, 5, 4, 5, 1, 5],
       [0, 4, 0, 3, 0, 3],
       [4, 2, 1, 2, 3, 0],
       [3, 1, 5, 1, 5, 2]], dtype=int32)
>>> 
>>> 
>>> np.take_along_axis(a, ind, axis = 0)
array([[ 9, 15,  0, 61, 17, 23],
       [10,  0,  4, 38, 16,  8],
       [28, 16, 46, 62, 20, 54],
       [43, 54, 62, 62, 53, 55],
       [89, 74, 93, 80, 53, 78],
       [88, 92, 85, 97, 95, 85]])

Получить разбиение другого массива на основе полученного индекса так же просто:

>>> b = np.vstack((np.zeros((3, 6)), np.ones((3, 6))))
>>> b
array([[0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0.],
       [1., 1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1., 1.]])
>>> 
>>> np.take_along_axis(b, ind, axis = 0)
array([[1., 1., 1., 1., 1., 0.],
       [0., 0., 0., 0., 0., 1.],
       [0., 1., 1., 1., 0., 1.],
       [0., 1., 0., 1., 0., 1.],
       [1., 0., 0., 0., 1., 0.],
       [1., 0., 1., 0., 1., 0.]])