Сахир и нубийский рынок

Бинпоиск по ответу + куча для поиска k наименьших

Есть n сувениров с базовыми стоимостями a[1], a[2], ..., a[n]. Если Сахир покупает k сувениров, то стоимость сувенира с номером i становится

costi = a[i] + i · k

То есть цена каждого сувенира зависит от того, сколько всего сувениров Сахир решил купить. Нужно найти максимальное k, при котором можно уложиться в бюджет S, и для этого k — минимальную возможную сумму.

Пусть f(k) — минимальная возможная сумма покупки ровно k сувениров. Утверждение: если f(k) > S, то и f(k+1) > S. Значит свойство «k сувениров помещается в бюджет» монотонно по k — можно делать бинарный поиск по ответу.

Остаётся научиться для фиксированного k быстро считать f(k). Это сводится к задаче: найти сумму k наименьших среди чисел a[i] + i·k для i = 1..n.

Классический приём: если нужны k наименьших элементов из потока, используем max-heap размера k. Держим в куче текущих кандидатов, а на вершине — самого «плохого» (наибольшего) из принятых. Когда приходит новый элемент:

• если в куче меньше k элементов — просто добавляем;
• иначе сравниваем с вершиной: если новый меньше вершины — вытесняем вершину, кладём нового;
• если новый больше или равен вершине — игнорируем, он точно не войдёт в k наименьших.

Параллельно поддерживаем sum — текущую сумму элементов в куче. После прохода по всем n элементам в куче лежат ровно k наименьших, а sum и есть f(k).

1. lo = 0, hi = n, ответ bestK = 0, bestT = 0.
2. На каждой итерации берём mid = (lo + hi) / 2, считаем cost = f(mid).
3. Если cost ≤ S — mid подходит, пробуем больше: bestK = mid, bestT = cost, lo = mid + 1.
4. Иначе hi = mid - 1.

1. Если k == 0 — сумма 0, сразу возвращаем.
2. Идём по i = 1..n, считаем cost = a[i] + i·k.
3. Пока в куче меньше k — пушим.
4. Иначе если cost < top — вытесняем top, пушим cost, обновляем sum.
5. В конце возвращаем sum.

Выбери массив a и зафиксированное k (внешний бинпоиск здесь опущен для наглядности). Слева — список сувениров со значениями a[i] + i·k, справа — текущее состояние max-heap. Каждый шаг обрабатывает одного кандидата.

Обрати внимание на момент, когда куча заполнилась и приходит новый кандидат с меньшей стоимостью — вершина (самый дорогой из принятых) вытесняется, сумма уменьшается. Те сувениры, которые игнорируются, становятся блёклыми.

#include <bits/stdc++.h>
using namespace std;

int n;
long long S;
vector<long long> a;

// Минимальная сумма покупки ровно k сувениров
long long minCost(int k) {
    if (k == 0) return 0;
    // max-heap: храним k наименьших значений a[i] + i*k
    priority_queue<long long> pq;
    long long sum = 0;
    for (int i = 1; i <= n; i++) {
        long long cost = a[i] + (long long)i * k;
        if ((int)pq.size() < k) {
            pq.push(cost);
            sum += cost;
        } else if (cost < pq.top()) {
            sum -= pq.top(); pq.pop();
            pq.push(cost);
            sum += cost;
        }
    }
    return sum;
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    cin >> n >> S;
    a.assign(n + 1, 0);
    for (int i = 1; i <= n; i++) cin >> a[i];

    // Бинпоиск: максимальное k, при котором minCost(k) <= S
    int lo = 0, hi = n, bestK = 0;
    long long bestT = 0;
    while (lo <= hi) {
        int mid = (lo + hi) / 2;
        long long cost = minCost(mid);
        if (cost <= S) {
            bestK = mid;
            bestT = cost;
            lo = mid + 1;
        } else {
            hi = mid - 1;
        }
    }

    cout << bestK << ' ' << bestT << endl;
    return 0;
}

import heapq
import sys

def min_cost(a, n, k):
    if k == 0:
        return 0
    # heapq — min-heap, а нам нужен max-heap,
    # поэтому кладём отрицательные значения
    heap = []
    s = 0
    for i in range(1, n + 1):
        cost = a[i] + i * k
        if len(heap) < k:
            heapq.heappush(heap, -cost)
            s += cost
        elif cost < -heap[0]:         # текущий максимум в куче = -heap[0]
            s -= -heapq.heapreplace(heap, -cost)
            s += cost
    return s

def main():
    data = sys.stdin.read().split()
    n, S = int(data[0]), int(data[1])
    a = [0] + [int(x) for x in data[2:2 + n]]

    lo, hi = 0, n
    best_k, best_t = 0, 0
    while lo <= hi:
        mid = (lo + hi) // 2
        cost = min_cost(a, n, mid)
        if cost <= S:
            best_k, best_t = mid, cost
            lo = mid + 1
        else:
            hi = mid - 1

    print(best_k, best_t)

main()

При n = 10⁵ это примерно 10⁵ · 17 · 17 ≈ 3·10⁷ операций — укладывается в 1-2 секунды с запасом.

Раньше мы решали эту задачу другим способом: фиксируем k, считаем все n значений a[i] + i·k, сортируем массив и берём сумму первых k. Сравним два подхода честно.

Формально одинаково. Разница — в константах и памяти.

• Память. Сортировка копирует весь массив data = nums на каждую проверку. Куча хранит только k элементов.
• Малое k. Когда бинпоиск пробует маленькое k (например, 10), сортировка всё равно возится со всем массивом, а куча делает n дешёвых сравнений с вершиной.
• Поток данных. Если данные приходят по одному — куча работает на лету, сортировка требует всё в памяти.

• Код проще. Буквально три строчки: сгенерировать, отсортировать, просуммировать первые k. Меньше шансов ошибиться.
• Кэш-локальность. std::sort на vector<long long> отлично ложится на кэш — последовательный доступ, предсказуемые ветвления. priority_queue на куче-дереве к кэшу менее дружелюбна. На реальных данных std::sort часто работает быстрее, несмотря на худшую асимптотику.

В том варианте бинпоиск стартовал с left = 1. Если даже один сувенир не по карману (пример 3 из условия: n=1, S=7, a=[7]), переменная best остаётся неинициализированной — на выходе мусор. Лечится либо left = 0, либо явной инициализацией best = {0, 0} перед циклом. В решении с кучей мы начинаем с lo = 0 и явно инициализируем bestK = 0, bestT = 0 — случай «ничего не купим» обрабатывается корректно.