Редакция для Сумма k наибольших
Submitting an official solution before solving the problem yourself is a bannable offence.
1. Идея
Для каждого запроса нужно узнать сумму k наибольших прибылей.
Если отвечать на каждый запрос отдельно, например каждый раз заново искать k лучших элементов, получится слишком медленно: запросов до 2 * 10^5.
Значит, нужно сделать подготовку один раз, а потом отвечать быстро.
Главная идея:
- отсортировать все прибыли по убыванию;
- посчитать массив префиксных сумм;
- тогда сумма
kнаибольших элементов будет просто равна сумме первыхkэлементов в отсортированном массиве.
2. Наблюдения
После сортировки по убыванию массив выглядит так:
a[0]— самая большая прибыль,a[1]— вторая по величине,- ...
a[k - 1]—k-я по величине.
Значит, сумма
kнаибольших прибылей — это сумма первыхkэлементов.Чтобы быстро получать сумму первых
kэлементов, удобно построить префиксные суммы:pref[0] = 0pref[1] = a[0]pref[2] = a[0] + a[1]- ...
pref[k]— сумма первыхkэлементов.
Тогда ответ на запрос
k— это простоpref[k].В задаче возможны отрицательные числа. Это никак не мешает:
- мы всё равно сортируем по убыванию;
- если
kбольшое, в сумму могут попасть отрицательные значения, и это нормально.
Сумма может не поместиться в 32-битный тип, поэтому в C++ нужен
long long.
3. Алгоритм
- Считать
n. - Считать массив прибылей
a. - Отсортировать
aпо убыванию. - Построить массив префиксных сумм
prefдлиныn + 1:pref[0] = 0pref[i + 1] = pref[i] + a[i]
- Считать число запросов
q. - Для каждого запроса:
- считать
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])
Комментарии