Частичные сортировки: секционирование

Иногда нам не требуется сортировать весь массив, а просто нужно найти K наименьших значений в нем. Библиотека NumPy предоставляет для этой цели функцию np.partition. Функция np.partition принимает на входе массив и число K. Результат представляет собой новый массив с K наименьшими значениями слева от точки разбиения и остальные значения справа от нее в произвольном порядке:

x = np.array([7, 2, 3, 1, 6, 5, 4])

np.partition(x, 3)

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

Аналогично сортировке можно секционировать по произвольной оси многомерного массива:

np.partition(X, 2, axis=1)

[2, 3, 4, 7, 6, 7],

[1, 2, 4, 5, 7, 7],

[0, 1, 4, 5, 9, 5]])

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