Редакция для Сумма k наибольших


Remember to use this editorial only when stuck, and not to copy-paste code from it. Please be respectful to the problem author and editorialist.
Submitting an official solution before solving the problem yourself is a bannable offence.

1. Идея

Для каждого запроса нужно узнать сумму k наибольших прибылей.

Если отвечать на каждый запрос отдельно, например каждый раз заново искать k лучших элементов, получится слишком медленно: запросов до 2 * 10^5.

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

Главная идея:

  • отсортировать все прибыли по убыванию;
  • посчитать массив префиксных сумм;
  • тогда сумма k наибольших элементов будет просто равна сумме первых k элементов в отсортированном массиве.

2. Наблюдения

  1. После сортировки по убыванию массив выглядит так:

    • a[0] — самая большая прибыль,
    • a[1] — вторая по величине,
    • ...
    • a[k - 1]k-я по величине.

    Значит, сумма k наибольших прибылей — это сумма первых k элементов.

  2. Чтобы быстро получать сумму первых k элементов, удобно построить префиксные суммы:

    • pref[0] = 0
    • pref[1] = a[0]
    • pref[2] = a[0] + a[1]
    • ...
    • pref[k] — сумма первых k элементов.
  3. Тогда ответ на запрос k — это просто pref[k].

  4. В задаче возможны отрицательные числа. Это никак не мешает:

    • мы всё равно сортируем по убыванию;
    • если k большое, в сумму могут попасть отрицательные значения, и это нормально.
  5. Сумма может не поместиться в 32-битный тип, поэтому в C++ нужен long long.

3. Алгоритм

  1. Считать n.
  2. Считать массив прибылей a.
  3. Отсортировать a по убыванию.
  4. Построить массив префиксных сумм pref длины n + 1:
    • pref[0] = 0
    • pref[i + 1] = pref[i] + a[i]
  5. Считать число запросов q.
  6. Для каждого запроса:
    • считать k;
    • вывести pref[k].

4. Почему это работает

После сортировки по убыванию первые k элементов массива — это ровно k наибольших прибылей из всех имеющихся.

Значит, ответ на запрос равен:

  • a[0] + a[1] + ... + a[k - 1]

По определению префиксных сумм:

  • pref[k] = a[0] + a[1] + ... + a[k - 1]

Следовательно, для каждого запроса ответом является pref[k].

Так как это верно для любого допустимого k, алгоритм корректно отвечает на все запросы.

5. Сложность

  • Сортировка: O(n log n)
  • Построение префиксных сумм: O(n)
  • Ответ на каждый запрос: O(1)
  • Все запросы вместе: O(q)

Итоговая сложность:

  • O(n log n + q)

Память:

  • O(n) на массив и префиксные суммы.

6. Код на C++17

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

int main() {
    int n;
    cin >> n;

    vector<long long> a(n);
    for (int i = 0; i < n; i++) {
        cin >> a[i];
    }

    sort(a.begin(), a.end(), greater<long long>());

    vector<long long> pref(n + 1, 0);
    for (int i = 0; i < n; i++) {
        pref[i + 1] = pref[i] + a[i];
    }

    int q;
    cin >> q;
    for (int i = 0; i < q; i++) {
        int k;
        cin >> k;
        cout << pref[k] << '\n';
    }

    return 0;
}

7. Код на Python 3

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

a.sort(reverse=True)

pref = [0] * (n + 1)
for i in range(n):
    pref[i + 1] = pref[i] + a[i]

q = int(input())
for _ in range(q):
    k = int(input())
    print(pref[k])

Комментарии

Еще нет ни одного комментария.