Метод partition

Метод partition выполняет частичную сортировку массива вдоль заданной оси. Он гарантирует, что элемент с индексом k окажется на своей правильной позиции, как если бы массив был полностью отсортирован. Все элементы слева от него будут меньше или равны, а все элементы справа - больше или равны. Это полезно, когда нужно найти несколько наименьших или наибольших значений без полной сортировки.

Параметр kth задает индекс или список индексов, которые должны оказаться на правильных позициях. Параметр axis определяет ось, вдоль которой выполняется частичная сортировка. Метод работает на месте, но если передать параметр kind, можно задать алгоритм сортировки.

Синтаксис

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

Пример

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

import numpy as np

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

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

[2 1 3 7 9 5 8 4 6]

Элемент с индексом 2 стал равен 3. Все элементы слева (2, 1) меньше или равны 3, а справа - больше или равны 3.

Пример

Выполним частичную сортировку по нескольким индексам одновременно:

import numpy as np

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

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

[2 1 3 4 5 6 9 7 8]

Индексы 2 и 5 заняли правильные позиции (3 и 6 соответственно). Элементы слева от каждого из них меньше или равны, справа - больше или равны.

Пример

Применим partition к двумерному массиву вдоль строк:

import numpy as np

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

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

[[4 0 2 1]
 [5 3 7 3]
 [9 8 8 6]]

Вдоль оси 0 (по строкам) элемент с индексом 1 в каждом столбце оказался на правильной позиции. В первом столбце значения стали (4, 5, 9), во втором (0, 3, 8) и так далее.

Пример

Выполним частичную сортировку вдоль столбцов двумерного массива:

import numpy as np

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

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

[[1 3 9 7]
 [2 5 8 6]
 [0 3 8 4]]

В каждой строке элемент с индексом 1 оказался на правильной позиции. Первая строка стала (1, 3, 9, 7), вторая - (2, 5, 8, 6), третья - (0, 3, 8, 4).

Пример

Найдем три наименьших элемента в массиве с помощью частичной сортировки:

import numpy as np

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

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

[2 1 3]

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

Пример

Найдем три наибольших элемента, используя отрицательный индекс:

import numpy as np

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

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

[6 9 8]

Последние три элемента содержат три наибольших значения. Они не упорядочены, но все они больше или равны остальным элементам.

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

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