EduBrick

Готовый поиск в C++ и Python

lower_bound, upper_bound, equal_range и модуль bisect. Что они возвращают и где их применять нельзя.

3 мин

Писать бинарный поиск руками нужно, когда ищется граница по предикату. Для поиска в отсортированном контейнере всё уже написано.

C++

auto it = lower_bound(a.begin(), a.end(), x);   // первый элемент >= x
auto jt = upper_bound(a.begin(), a.end(), x);   // первый элемент > x
bool has = binary_search(a.begin(), a.end(), x);
auto [from, to] = equal_range(a.begin(), a.end(), x);  // сразу обе границы

Все возвращают итераторы, а не индексы. Чтобы получить индекс, вычитают a.begin():

int index = lower_bound(a.begin(), a.end(), x) - a.begin();
int count = upper_bound(a.begin(), a.end(), x) - lower_bound(a.begin(), a.end(), x);

Если элемента нет, lower_bound возвращает a.end() — итератор за концом. Разыменовывать его нельзя, сравнивать с a.end() нужно обязательно.

С компаратором

Третий аргумент задаёт порядок, и он должен совпадать с тем, по которому массив отсортирован:

// массив отсортирован по убыванию
auto it = lower_bound(a.begin(), a.end(), x, greater<int>());

Забыть компаратор на убывающем массиве — тихая ошибка: функция ничего не заметит и вернёт мусор.

Где применять нельзя

lower_bound на std::set и std::map работает, но вызывать надо метод контейнера, а не свободную функцию:

s.lower_bound(x);                       // правильно: логарифм
lower_bound(s.begin(), s.end(), x);     // неправильно: линия

Разница в том, что свободная функция умеет прыгать по индексам только у итераторов с произвольным доступом. У дерева таких итераторов нет, и она честно идёт по элементам подряд — получается O(n)O(n) вместо O(logn)O(\log n). Компилятор об этом не предупредит.

По той же причине бессмысленно звать её на std::list.

Python

from bisect import bisect_left, bisect_right, insort

i = bisect_left(a, x)     # первая позиция, куда можно вставить x
j = bisect_right(a, x)    # последняя такая позиция
count = j - i             # сколько раз x встречается
has = i < len(a) and a[i] == x

Здесь возвращаются индексы, а не итераторы, и это удобнее.

У обеих функций есть параметры lo и hi — поиск в части списка. А начиная с Python 3.10 есть и key=, что избавляет от построения отдельного массива ключей:

i = bisect_left(people, 30, key=lambda p: p.age)

insort и его ловушка

insort(a, x) вставляет элемент с сохранением порядка. Позиция ищется за логарифм, но сама вставка сдвигает хвост списка и стоит O(n)O(n).

Для поддержания отсортированного набора из многих элементов это плохой выбор: получится O(n2)O(n^2). Если набор большой и меняется часто, нужна другая структура — куча, если хватает минимума, или отсортированный контейнер из сторонней библиотеки.

Своё или готовое

Готовое — когда ищете элемент в отсортированном контейнере. Своё — когда ищете границу по предикату: библиотечные функции про массивы, а поиск по ответу к массиву не сводится.

Промежуточный случай: partition_point в C++ ищет границу ровно по предикату и делает то же, что ручной поиск:

auto it = partition_point(a.begin(), a.end(), [&](int value) { return value < x; });

Требование одно — предикат должен быть монотонным по диапазону, то есть сначала везде истинным, потом везде ложным.