問題概要

$N$ 体の敵がいて、敵に付随するスコアはそれぞれ $C_{1}, \dots, C_{N}$ で与えられる (負数になることもある)。

これらの敵を順に倒していきたい。Boss Score, Total Score と呼ばれる値が初期状態ではともに 0 となる。敵 $i$ を倒すとき、次のような挙動になる。

Total Score に現在の Boss Score を加算する
Boss Score に $C_{i}$ を加算する

また任意のタイミングで魔法を唱えることができて、魔法を唱えると Boss Score の値を 0 にすることができる。ただし魔法は $K$ 回までしか使えない。

$N$ 体の敵をすべて倒し終えたときの、Total Score の最大値を求めよ。

制約

$N, K \le 5 \times 10^{5}$

解法 (1)：僕の解法

状況が複雑なので、上手に言い換えよう。そうすると、次のように言い換えることができる。

まず $C$ を適切に並び替える
それらを $K+1$ 個以下の連続区間に分割する
各区間の要素を $a_{0}, \dots, a_{L-1}$ としたとき、その区間の敵を倒すことでえられるスコアは $0 \times a_{0} + 1 \times a_{1} + 2 \times a_{2} + \dots + (L-1) \times a_{L-1}$ となる

つまり、各区間において、各要素の寄与率が $0, 1, 2, \dots$ となるものと考えられる。これでだいぶわかりやすくなったけど、さらに「寄与率」ごとに要素を並び替えることにすると、次のように言い換えられる。

$C$ を適切に並び替える
$C$ を連続する区間ごとに分割していく、ただし各区間の長さは $K+1$ 以下でなければならない
$i$ 番目の区間の寄与率は $i$ となる (0-indexed)

このノリでサンプル 3 を解釈すると、次のようになって解が 71 になる。

値	-9	-9	-8	-6	-5	-5	-3	-2	1	1	3	4	5
寄与率	0	0	0	1	1	2	3	4	5	6	7	8	9

解法

ここまでわかれば、寄与率をどう分配するかの問題になる。

前半は、寄与率の増加を抑えた方がいい ( $K$ 個進むごとに寄与率を 1 上げる)
後半は、寄与率を積極的に上昇させた方がいい (1 個進むごとに寄与率を 1 上げる)

という絶妙なバランスになっている。寄与率は

(0, 0, ..., 0, 1, 1, ..., 1)

のようなベクトルの総和とみなすことができる。このようなベクトルを足した方がよいかどうかを考察してみる。具体的には

区間 [v, N) の総和が 0 以上ならば、v 番目から先の寄与率を 1 ずつ増加させた方がよい
区間 [v, N) の総和が 0 未満ならば、v 番目から先の寄与率をなるべく小さくしたい

というふうにすれば OK。計算量はソートがボトルネックで $O(N \log N)$ となる。

コード

#include <bits/stdc++.h>
using namespace std;
using pint = pair<int, int>;
using pll = pair<long long, long long>;
template<class T> inline bool chmax(T& a, T b) { 
    if (a < b) { a = b; return 1; } 
    return 0; 
}

int main() {    
    long long N, K;
    cin >> N >> K;
    vector<long long> C(N);
    long long tot = 0;
    for (int i = 0; i < N; ++i) cin >> C[i], tot += C[i];
    sort(C.begin(), C.end());

    int div = 0;
    long long sum = 0;
    for (int i = 0; i < N; ++i) {
        sum += C[i];
        if (tot - sum < 0) div = i+2;
    }

    long long res = 0;
    long long cur = -1, num = 0;
    for (int i = 0; i < div; ++i) {
        if (num == 0) ++cur;
        res += C[i] * cur;
        num = (num + 1) % (K+1);
    }
    for (int i = div; i < N; ++i) {
        ++cur;
        res += C[i] * cur;
    }
    cout << res << endl;
}

解法 (2)：priority_queue を用いた貪欲

言い換えるのではなく、元の Boss Score と Total Score について考える。ただし魔法をとなえる代わりに、 $K$ 系列に分解して、それぞれについて最適な敵を割り当てていく考え方をする。

ここで敵を倒す順序としては、各系列については「 $C$ の大きい順」にするのがよいことはわかる。そこで、あらかじめ $C$ を大きい順にソートしておく。そして各 $C_{i}$ について、どの系列で倒していくかを Greedy に決めていけば OK。

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

int main() {    
    long long N, K;
    cin >> N >> K;
    vector<long long> C(N);
    for (int i = 0; i < N; ++i) cin >> C[i];
    sort(C.begin(), C.end(), greater<long long>());

    long long res = 0;
    priority_queue<long long> que;

    // 1 体目を倒すときの Boss Score は 0
    for (int k = 0; k <= K; ++k) que.push(0);
    for (int i = 0; i < N; ++i) {
        // K 系列のうちの Boss Score の最大値
        long long tmp = que.top(); 
        que.pop();

        // Total Score に tmp を加算して、tmp に敵 C[i] を加算
        res += tmp;
        tmp += C[i];

        // 改めて que に push
        que.push(tmp);
    }
    cout << res << endl;
}