けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

Codeforces Round #687 (Div. 1) C. New Game Plus! (R2200)

すごい面白い!!

問題概要

 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;
}