searchsorted и partition — NumPy

Коротко

np.searchsorted(a, v) — бинарный поиск: индекс, куда вставить v, чтобы a остался отсортированным. Работает за O(log n).
Массив a должен быть отсортирован по возрастанию — это никто не проверяет.
side="left" (по умолчанию) ставит перед равными, side="right" — после. Разница между ними — число вхождений v.
v может быть массивом: ответ для каждого элемента сразу. Так значения раскладывают по интервалам.
np.partition(a, k) ставит на место k-й по величине элемент: слева меньшие, справа большие, но без сортировки внутри частей. Это быстрее полной сортировки, если нужны только top-k.

Синтаксис

np.searchsorted(a, v, side="left")
np.partition(a, k)           # k-й элемент на своём месте
np.partition(a, -k)[-k:]     # k наибольших (в произвольном порядке)
np.argpartition(a, k)        # то же, но индексы

Параметры

ПараметрЧто делаетПо умолчанию

aотсортированный по возрастанию одномерный массивобязательный

vзначение или массив значений для поискаобязательный

side"left" — перед равными, "right" — после"left"

kиндекс элемента, который встанет на своё место; отрицательный — с концаобязательный

Примеры

Данные для примеровиспользуются во всех примерах ниже

limits = np.array([0, 1000, 5000, 20000])      # границы тарифов, по возрастанию
scores = np.array([72, 45, 90, 58, 81, 39, 66]) # баллы участников

Позиция для вставки

np.searchsorted(limits, 3000)                  # куда вставить, чтобы порядок сохранился

Вывод

np.int64(2)

Равное значение: left и right

print(np.searchsorted(limits, 5000))                  # слева от равного
np.searchsorted(limits, 5000, side="right")           # справа от равного

Вывод

2
np.int64(3)

Номер интервала для каждого значения

spend = np.array([500, 1000, 7500, 25000])
np.searchsorted(limits, spend, side="right") - 1      # номер тарифа для каждой суммы

Вывод

array([0, 1, 2, 3])

Вставить с сохранением порядка

sorted_prices = np.array([90, 120, 250, 480])
pos = np.searchsorted(sorted_prices, 200)
np.insert(sorted_prices, pos, 200)

Вывод

array([ 90, 120, 200, 250, 480])

k наименьших и наибольших

part = np.partition(scores, 2)
print(part[2])                                 # на месте 2 — третий по величине
print(np.sort(part[:2]), np.sort(part[3:]))    # слева меньшие, справа большие; порядок внутри частей любой
np.sort(np.partition(scores, -3)[-3:])         # 3 наибольших (отсортированы уже для вывода)

Вывод

58
[39 45] [66 72 81 90]
array([72, 81, 90])

Индексы top-3 по убыванию

idx = np.argpartition(scores, -3)[-3:]         # индексы трёх лучших…
idx[np.argsort(scores[idx])[::-1]]             # …упорядоченные по убыванию баллов

Вывод

array([2, 4, 0])

Подводные камни

Неотсортированный массив

searchsorted не проверяет порядок и возвращает какое-то число без ошибки. Если данные не отсортированы, сначала np.sort или используйте argsort + sorter=.

Поиск в неотсортированном массиве

np.searchsorted(np.array([5, 1, 3]), 2)        # массив не отсортирован — ответ бессмысленный

Вывод

np.int64(2)

partition не сортирует части

После partition упорядочен только один элемент — k-й. Если top-k нужны по порядку, отсортируйте эти k элементов отдельно, как в примере выше.