Бинарный поиск по массиву

Список из двухсот тысяч чисел упорядочен, и нужно ответить на сто тысяч вопросов «есть ли здесь такое число».

Перебором это $2 \cdot 10^5 \cdot 10^5 = 2 \cdot 10^{10}$ действий — часы работы. Но упорядоченность даёт больше, чем кажется: сравнив запрос с одним элементом в середине, мы сразу узнаём, в какой половине искать. Вторую половину можно не смотреть никогда.

Каждое сравнение уменьшает область поиска вдвое. Из двухсот тысяч остаётся сто тысяч, потом пятьдесят, потом двадцать пять… За восемнадцать шагов остаётся один элемент. Восемнадцать вместо двухсот тысяч.

Миллиард — тридцать шагов. Это и есть $\log_2 n$ из прошлого занятия про сложность.

Упорядоченность обязательна. Без неё сравнение с серединой не говорит ничего: подходящее число может оказаться где угодно. Если данные пришли в произвольном порядке, их сначала сортируют — это $n \log n$, и при многих запросах такая плата окупается сразу.

Бинарный поиск обычно запоминают как заклинание — и почти всегда путают, где +1, где <=, а где mid. Заклинание запоминать не надо. Надо сформулировать, что мы знаем про границы, и код напишется сам.

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

• left — место, где условие заведомо не выполнено;
• right — место, где условие заведомо выполнено.

Это и есть инвариант: утверждение, верное перед каждым шагом и после него.

left = -1          # заведомо "нет": до начала списка
right = n          # заведомо "да": за концом списка

while right - left > 1:
    mid = (left + right) // 2
    if a[mid] < x:      # условие "уже не меньше x" не выполнено
        left = mid
    else:
        right = mid

print(right)       # первое место, где условие выполнено

Разберём, почему это верно.

Начальные значения. left = -1 — за левым краем условие не выполнено по договорённости, там ничего нет. right = n — за правым краем условие «выполнено», тоже по договорённости. Эти два места существуют не в списке, а в рассуждении, и обращаться к ним по индексу мы не будем.

Каждый шаг. mid строго между границами, потому что расстояние больше единицы. Куда бы ни пошёл ответ, инвариант сохраняется.

Конец. Цикл закончился, когда right - left = 1. Значит границы стоят вплотную: слева условие не выполнено, справа выполнено. Ровно на этой стыке и находится «первое место, где условие выполнено» — это right.

Обратите внимание: никакой отдельной проверки «а нашли ли мы» внутри цикла нет. Цикл всегда доводит границы до стыка, и ответ читается из инварианта, а не из удачного попадания.

И right = n — не описка. Ответ «подходящего элемента нет» тоже должен где-то помещаться, и он помещается как раз в позицию $n$.

Цикл while right - left > 1 закончился.

Что известно про границы?

Когда в списке есть повторы, «найти x» — вопрос неточный. Точных вопросов два, и оба решаются одним и тем же поиском с разным условием.

Первое место, где значение не меньше x. Условие: a[mid] >= x. Это начало куска из иксов, если они есть.

Первое место, где значение строго больше x. Условие: a[mid] > x. Это место сразу за концом куска.

Из этих двух чисел получается всё остальное:

Последняя строка — рабочая лошадка. Почти любой вопрос «сколько элементов со значением от и до» сводится к двум поискам и одному вычитанию.

Проверка на существование тоже выражается через первый поиск, но аккуратно:

pos = первое_не_меньше(x)
if pos < n and a[pos] == x:
    print("есть")

Порядок условий важен. Сначала pos < n, потом обращение a[pos] — иначе на запросе больше всех элементов программа упадёт с выходом за границы.

В Python оба поиска уже написаны — в модуле bisect.

from bisect import bisect_left, bisect_right

a = [1, 3, 3, 3, 7]

bisect_left(a, 3)    # 1 — первое место, где значение не меньше 3
bisect_right(a, 3)   # 4 — первое место, где значение больше 3
bisect_right(a, 3) - bisect_left(a, 3)   # 3 — столько троек

Названия говорят о том, куда встало бы новое значение: bisect_left — левее всех равных, bisect_right — правее всех равных. Если равных нет, оба возвращают одно и то же место.

У обеих функций есть необязательные границы lo и hi — поиск только в части списка:

bisect_left(a, 3, 2)     # искать начиная с места 2

На олимпиаде bisect использовать можно и нужно: он написан на C и работает заметно быстрее ручного цикла. Но писать поиск руками вы всё равно должны уметь — на следующем занятии условие перестанет быть «значение не меньше x», и готовой функции для него не окажется.

Ещё в модуле есть insort — вставка с сохранением порядка. Сама вставка сдвигает хвост списка, то есть стоит $O(n)$, и на больших списках она не спасает.

Список: 2 4 4 6 8 9 11.

Сколько в нём элементов, строго больших 4 и строго меньших 11? Введите целое число.

Правая граница равна n - 1. Самая частая. Пока ответ есть в списке, всё работает; как только подходящего элемента нет, вернётся последнее место вместо «за концом». Признак: неверные ответы ровно на запросах, больших всех элементов.

Бесконечный цикл. Возникает, когда границу двигают в mid, но условие цикла позволяет mid совпасть с границей. При left = mid и while left < right середина может остаться на месте — и программа зависнет. Форма while right - left > 1 от этого защищена: mid строго между границами.

Обращение к a[pos] без проверки. pos может оказаться равным $n$. Сначала проверяем номер, потом читаем значение.

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

Перепутанные строгие и нестрогие. «Не меньше» и «больше» отличаются одним знаком в условии, а ответы расходятся ровно на количество равных элементов. Если список без повторов, ошибку не видно — и она всплывает на тесте с повторами.

Как проверять себя

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

Лестница прежняя: пять разминочных, семь основных, три со звёздочкой. Зачёт при десяти решённых.

В части задач список дан по неубыванию, в части — в произвольном порядке. Читайте условие: сортировка на вашей совести.

И держите в голове формулировку через инвариант. Когда задача сложнее, чем «найти число», шаблон подвести может, а инвариант — нет.

Задача: «дан список из $n$ чисел по неубыванию; для каждого запроса $x$ вывести номер первого элемента, не меньшего $x$, а если такого нет — вывести $n + 1$». Нумерация с единицы.

Ученик написал:

n, x = map(int, input().split())
a = list(map(int, input().split()))

left = 0
right = n - 1
while left < right:
    mid = (left + right) // 2
    if a[mid] < x:
        left = mid + 1
    else:
        right = mid

print(left + 1)

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