Japan Alumni Group Summer Camp 2014 Day 4 D - 夕食

問題文

これ

概要

原文が日本語なので省略.

解法

ぱっと見ると, 日数とその日までで何回自炊したかというのを状態とする DP が出来そうだが, 今回は $n \leqq 5 \times 10^5$ なのでそれは厳しそう.
ここで, DP のキーにもなっている「自炊した回数」に着目する.

毎日自炊をする意識の高い学生である場合の日々の幸福度を表にすると以下の様になる.

1日目	2日目	3日目	$\cdots$	$i-1$日目	$i$日目	$i+1$日目	$\cdots$	$n$ 日目	合計
$PQ$	$P(Q+1)$	$P(Q+2)$	$\cdots$	$P(Q+i-2)$	$P(Q+i-1)$	$P(Q+i)$	$\cdots$	$P(Q+n-1)$	$P(nQ + \frac{n(n-1)}{2})$

ここで $i$ 日目に意志薄弱になって食堂に行ったとすると, 幸福度は次のように変化する.

1日目	2日目	3日目	$\cdots$	$i-1$日目	$i$日目	$i+1$日目	$\cdots$	$n$ 日目	合計
$PQ$	$P(Q+1)$	$P(Q+2)$	$\cdots$	$P(Q+i-2)$	$C_i$	$P(Q+i-2)$	$\cdots$	$P(Q+n-3)$	$P(nQ + \frac{n(n-1)}{2}) + C_i - P(Q + i - 1) - 2P(n - i)$

食堂に行くことにした式と行く前の式の差分を取ると, 差分は $(C_i + Pi) + P - P(Q + 2n)$ となる. つまり, 1 日だけ食堂に行くのであれば, $C_i + Pi$ が最も大きい物について確かめるだけで良い, ということが分かる.

同様に表を書いていき考察することで, $k$ 日間食堂に行くときの差分は $\left(\displaystyle\sum_{i=1}^{k}C_{a_i}+P a_{i}\right) + k^2P - kP(Q + 2n)$ となる. これは $C_{a_i} + P a_{i}$ のみに依存する値となるから, $C_{a_i} + P a_{i}$ の降順に足し上げて, 毎日自炊をする場合との差分を確かめてやれば良い.

コード

#include <bits/stdc++.h>

using namespace std;

int main()
{
    long long n, p, q;
    vector<long long> v;
    
    scanf("%lld %lld %lld", &n, &p, &q);
    
    for (int i = 1; i <= n; i++){
        long long x;
        scanf("%lld", &x);
        v.push_back(x + p * i);
    }
    
    sort(v.rbegin(), v.rend());
    
    long long diff = 0;
    long long s = 0;
    for (int i = 0; i < n; i++){
        s += v[i];
        diff = max(diff, s + (long long)(i + 1) * (i + 1) * p - (long long)(i + 1) * p * (q + 2 * n));
    }
    
    printf("%lld\n", p * (n * q + n * (n - 1) / 2) + diff);
    
    return (0);
}