Подсчёты и запросы бинарным поиском

Бинарный поиск чаще используют не для «найти элемент», а для «посчитать, сколько». Приём один и тот же: находим две границы и вычитаем.

Сколько раз встречается

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

lower_bound даёт первое вхождение, upper_bound — позицию сразу за последним. Разность — количество. Если элемента нет, оба итератора совпадут и разность будет нулём — отдельная проверка не нужна.

В Python то же самое:

from bisect import bisect_left, bisect_right
count = bisect_right(a, x) - bisect_left(a, x)

Сколько элементов в диапазоне

Тот же приём для полуинтервала $[l, r]$:

int count = upper_bound(a.begin(), a.end(), r) - lower_bound(a.begin(), a.end(), l);

Это основа целого класса решений: отсортировали массив один раз, и дальше каждый запрос «сколько чисел в промежутке» стоит логарифм вместо прохода.

Ближайшее число

Задача: найти в отсортированном массиве элемент, ближайший к $x$.

Кандидатов ровно два: первый элемент не меньше $x$ и предыдущий. Всё, что дальше, заведомо хуже.

int nearest(const vector<int>& a, int x) {
    auto it = lower_bound(a.begin(), a.end(), x);
    int best = INT_MAX, answer = 0;
    if (it != a.end())     { best = abs(*it - x); answer = *it; }
    if (it != a.begin()) {
        int prev = *(it - 1);
        if (abs(prev - x) < best) answer = prev;
    }
    return answer;
}

Обе проверки обязательны. it == a.end() бывает, когда $x$ больше всех; it == a.begin() — когда $x$ меньше всех. Разыменование в этих случаях — обращение за границу массива, и «работает же на сэмплах» здесь ничего не значит.

Отдельно отметим: делать два независимых бинарных поиска — за первым большим и за последним меньшим — не нужно. Один lower_bound плюс шаг назад даёт то же самое вдвое дешевле.

Встречается ли x на отрезке индексов

Дан массив (не отсортированный) и запросы: встречается ли значение $x$ среди $a_l \dots a_r$?

Сортировать нельзя — индексы важны. Но можно построить обратный указатель: для каждого значения — отсортированный список позиций, где оно стоит.

unordered_map<int, vector<int>> positions;
for (int i = 0; i < n; i++) positions[a[i]].push_back(i);
// внутри каждого вектора индексы уже возрастают

bool occurs(int x, int l, int r) {
    auto found = positions.find(x);
    if (found == positions.end()) return false;
    const vector<int>& p = found->second;
    auto it = lower_bound(p.begin(), p.end(), l);
    return it != p.end() && *it <= r;
}

Первое вхождение начиная с $l$ — и если оно не дальше $r$, ответ положительный. Построение — $O(n)$, запрос — $O(\log n)$.

Тем же способом считается количество вхождений на отрезке: upper_bound(p, r) - lower_bound(p, l).

Медиана объединения двух массивов

Даны два отсортированных массива. Найти медиану их объединения, не сливая их, — за логарифм.

Идея: провести разрез. Пусть из первого массива в левую половину уходит $i$ элементов, тогда из второго — ровно $\frac{n+m+1}{2} - i$. Разрез корректен, когда всё слева не больше всего справа:

$$a_{i-1} \le b_j \quad\text{и}\quad b_{j-1} \le a_i$$

Бинарим по $i$.

long long median(const vector<long long>& a, const vector<long long>& b) {
    if (a.size() > b.size()) return median(b, a);   // бинарим по короткому
    int n = a.size(), m = b.size(), half = (n + m + 1) / 2;
    int lo = 0, hi = n;
    while (lo <= hi) {
        int i = (lo + hi) / 2, j = half - i;
        long long aLeft  = i > 0 ? a[i - 1] : LLONG_MIN;
        long long aRight = i < n ? a[i]     : LLONG_MAX;
        long long bLeft  = j > 0 ? b[j - 1] : LLONG_MIN;
        long long bRight = j < m ? b[j]     : LLONG_MAX;
        if (aLeft <= bRight && bLeft <= aRight) return max(aLeft, bLeft);
        if (aLeft > bRight) hi = i - 1; else lo = i + 1;
    }
    return 0;
}

Три детали, без которых не работает:

• бинарим по короткому массиву, иначе $j$ вылетит за границы;
• пустые части закрываются бесконечностями LLONG_MIN и LLONG_MAX, что убирает все проверки на края;
• half считается с округлением вверх — тогда при нечётной сумме длин медиана оказывается слева, и брать нужно max левых.

Есть и более простой путь к тому же результату — бинарный поиск по значению медианы с проверкой «сколько элементов не превосходит $v$». Он даёт $O(\log(\max))$ вместо $O(\log n)$ и пишется втрое короче; на практике почти всегда выбирают его.

MEX после k операций

Дано множество натуральных чисел. Операция: добавить к нему его собственный MEX — наименьшее натуральное, которого в множестве нет. Что окажется добавлено на $k$-м шаге?

Каждая операция затыкает ровно одну «дыру» в натуральном ряду, слева направо. Значит, ответ — $k$-я по счёту дыра.

Бинарим по границе $M$: сколько дыр среди чисел от 1 до $M$? Это $M$ минус количество элементов множества, не превосходящих $M$, — а второе как раз считается upper_bound по отсортированному массиву.

int holes = M - (upper_bound(a.begin(), a.end(), M) - a.begin());

Ищем минимальное $M$, где число дыр достигает $k$. Функция дыр неубывающая, значит бинарный поиск применим.

При ограничениях до $10^5$ можно и просто пройтись по ряду. Но при значениях до $10^9$ прохода нет, а бинарный поиск работает без изменений — это его типичное преимущество.

Общий шаблон

Все задачи выше сводятся к одной формуле:

сколько объектов удовлетворяет условию = (граница справа) − (граница слева)

И к одному правилу: если считаете «сколько», ищите две границы, а не элемент. Поиск конкретного элемента почти всегда лишний шаг.