mayoko’s diary

プロコンとかいろいろ。

AOJ 2642: Dinner

解法

食堂全振りをベースに考えます。

ここから部分問題として, 「j 日目を自炊にしたらどうなるか」を考えてみます。
こうすると, 当然のことながら C[j] の幸福度はなくなります。また, j 日目までに自炊パワーは j 減っているので, かわりに幸福度は Q-j 追加されます。

ただ, これは自炊にする日が一日だけだった時の話で, 複数あるときはより複雑になります。ただ, 考えるベースは同じです。

「j 日目までに自炊パワーが j 減っている」というのは使いましょう。つまり, 「j 日目を選択するたびに追加される幸福度は Q-j」と考えます。ただし, すでに i 日間自炊していたとすると, 自炊パワーは 1 上がるので, 合計すると i*(i-1)/2 * (1-(-1)) = i*(i-1) だけ幸福度は上がっています。

以上により,

sum を C[0]+...+C[N-1] に初期化
for (int i in 0..N) {
  ans = max(ans, sum+P*Q*i+P*i*(i-1))
  sum -= P*id[i]
  sum -= C[id[i]]
}

のような擬似コードで, 「i 日自炊する」という前提のもとでの幸福度を求められます。ただ, 幸福度を最大化するには id[i] は順番通りにやれば良いとは限らないのでそこは工夫しなければいけません。

上のコードを見ればわかるように, C[id[i]]+P*id[i] の値が小さい順にソートすると, 「i 日間自炊する」という条件のもとで幸福度を最大にすることが出来ます。

const int MAXN = 500500;
int C[MAXN];

int main() {
    int N;
    ll P, Q;
    cin >> N >> P >> Q;
    ll sum = 0;
    vector<pair<ll, int> > ps;
    for (int i = 0; i < N; i++) {
        cin >> C[i];
        sum += C[i];
        ps.emplace_back(C[i]+P*i, i);
    }
    sort(ps.begin(), ps.end());
    ll ans = -1ll<<55;
    for (int i = 0; i <= N; i++) {
        ans = max(ans, sum+P*i*(i-1)+P*Q*i);
        if (i < N) {
            sum -= C[ps[i].second] + P*ps[i].second;
        }
    }
    cout << ans << endl;
    return 0;
}