Amusement Park [AtCoder Beginner Contest 216 E]

https://atcoder.jp/contests/abc216/tasks/abc216_e

解説

https://atcoder.jp/contests/abc216/submissions/25449841

この問題は貪欲法で解ける。
勉強熱心な方はDPを思いついたかもしれないが、たぶん正常な反応だと思う。
貪欲法で解けそうな問題が実際はDPでしたというパターンも、逆のパターンもかなり見てきた。

今回は貪欲法の方針を見つけるのはそれほど難しくなく、実装がちょっと大変。

貪欲法の方針

貪欲法で解きますという風な前提があれば、その方針を見つけることはそれほど難しくないだろう。
楽しさを最大化するには、「各操作時に楽しさが最大のアトラクションを選ぶことを繰り返す」という貪欲をすればいい。
その貪欲法で問題ないかという見積もりをする必要があるが、パッと考えた感じ反例が思い浮かばず、
かなり自明な貪欲法に思えたのでそのまま突っ切って実装してACした。
正直これで何回も痛い目を見ているのだが、貪欲法の見積もりって難しいですよね…
基本DPから考えるといいと思います。で、不可能になったら貪欲法という感じで。

実装

さて、少し回り道をしたが、「各操作時に楽しさが最大のアトラクションを選ぶことを繰り返す」で貪欲法を実践する。
言われたまま実装をするとKの値は大きいのでTLEしてしまう。
大きい数を減らしていくと、次に大きい数と一致するときが来る。
この間は地道に-1するのではなく一気に減らすような工夫をすると貪欲法を高速化できる。
このように一定の操作が行われる部分については計算で省略をしながら高速に貪欲法をやっていく。
Aの値は大きい順で使用するので、降順ソートした上で、先頭から使用していこう。

サンプル1の例で考えてみよう。
先頭が102で次が100なので、その間の102+101が使用されることになる。
その次は、2番目が100でその次が50なので、100+99+98+...+51が使用されることになる。
ただ、2回目は1番目が2番目と同じ値にまで減らされているので(100+99+98+...+51)×2が使用される感じになる。
なので、基本的には(A[i] + ... + (A[i + 1] + 1))×(i + 1)を足していけばいいことになる。
iは0始まり、0-indexedで表記している。

このように等差数列の和を利用するので、自作のライブラリtousa_sumを持ってきて計算した。
高校生で「等差数列の和」は習うのだが、競プロやる中学生は知ってそうだな…
まあ、忘れている社会人含めて「等差数列の和」でググれば式は出てくるのでそれを使って実装すればいい。
自分の実装が気になるなら、提出URLに飛んで見てみてほしい。

半端な場合はどうする？

基本的には上記の方針でKを減らしていけばいいが、いつか中途半端になってしまう場合がある。
この場合には、これまでに数が揃っている個数分のA[i]が何回分同時に-1できるかというのを特定する。
これはKを(i+1)で割った答えdになる。
d回分は等差数列の和で計算する。
Kを(i+1)で割った余りmが、残った回数になるので、あとはその分だけ足し合わせて答えが出来上がる。

ll N, K, A[201010];
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N >> K;
    rep(i, 0, N) cin >> A[i];
    sort(A, A + N, greater<ll>());

    ll ans = 0;
    rep(i, 0, N) {
        ll diff = A[i] - A[i + 1];

        ll cnt = 1LL * diff * (i + 1);
        if (cnt <= K) {
            K -= cnt;
            ans += tousa_sum(A[i], -1LL, diff) * (i + 1);
        }
        else {
            ll d = K / (i + 1);
            ll m = K % (i + 1);
            ans += tousa_sum(A[i], -1LL, d) * (i + 1);
            ans += (A[i] - d) * m;
            K = 0;
        }
    }
    cout << ans << endl;
}