AOJ2642: Dinner(夕食)の非天才解法

少し前に Twitter で話題になった問題です.
AOJ: Dinner - AIZU ONLINE JUDGE
AtCoder: Japan Alumni Group Summer Camp 2014 Day 4: D - 夕食 | AtCoder

典型度は低めで実装も軽く,面白いと評判の問題です!

この問題は天才要素が強いとも言われているようですが,我ながらかなり自然に(=非天才的に)解けたと思うので,解説を書きます.

考察過程

  • DP か? → 制約大きいので無理
  • じゃあまずは数理最適化問題として定式化してみるか → 式変形したら解けた

という感じで,スムーズでした.

解法

数理最適化問題としての定式化

問題文に従って 1-indexed で書きます.

変数 $x_1,\dots,x_N$ を準備し,$i$ 日目に自炊するなら $x_i = 1$,自炊しないなら $x_i = 0$ と定めます.$x_1,\dots,x_N$ を用いて,最大化すべき目的関数を表したいです.

まず,$i$ 日目の自炊パワーは \[Q + \sum_{j=1}^{i-1} (2x_j - 1)\]と書けます.これを用いると,$N$ 日間で得られる幸福度の総和は\begin{equation}\label{obj}\red{\sum_{i=1}^N P \bigg( Q + \sum_{j=1}^{i-1} (2x_j - 1) \bigg) x_i+ \sum_{i=1}^N C_i(1-x_i)}\end{equation}と表せます.第 $1$ 項が自炊で得られる幸福度の総和で,第 $2$ 項が食堂で得られる幸福度の総和です.

これで目的関数が $x_1,\dots,x_N$ の式で書けました!

目的関数の整理

目的関数 \eqref{obj} は少しゴチャゴチャしていて見にくいので,$x_1,\dots,x_N$ の次数ごとに整理しましょう.$0$ 次の項から並べると,\[\sum_{i=1}^N C_i+ \sum_{i=1}^N \big( P(Q-i+1) - C_i \big) x_i+ \under{P \sum_{1 \leq j < i \leq N} 2x_i x_j} \] となります.

$1$ 次の項までしかなかったら簡単に最大化できるのですが,$2$ 次の項(赤線部)があるのが少し厄介です.そこで,\[\sum_{1 \leq j < i \leq N} 2x_i x_j = \bigg( \sum_{i=1}^N x_i \bigg)^2 - \sum_{i=1}^N x_i^2 \] という定番の等式を使ってみます.

さらに,$x_i = 0$ or $1$ なので $x_i^2 = x_i$ であることも用いると,目的関数は\begin{equation}\label{obj2}\red{\sum_{i=1}^N C_i+ \sum_{i=1}^N \big( P(Q-i) - C_i \big) x_i+ P\bigg( \sum_{i=1}^N x_i \bigg)^2}\end{equation}と変形できます.

貪欲法へ

ここまで来たらあとは簡単です.$S := \sum_{i=1}^N x_i$ とおいて $S$ を固定します.各 $S = 0,1,\dots,N$ に対し,式 \eqref{obj2} を $x_1,\dots,x_N$ について最大化すればよいです.

$S$ を固定したとき,$P(Q-i) - C_i$ が大きい $i$ から順に $x_i = 1$ とするのがよいので,事前にソートしておけば,全ての $S = 0,1,\dots,N$ について最大値が高速に求まり,答えを得ることができます.

時間計算量はソートがボトルネックとなり $O(N\log N)$ です.

AC コード

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < (n); ++i)
using namespace std;
using ll = long long;

int main() {
  ll n, p, q; cin >> n >> p >> q;

  ll sum0 = 0;
  vector<ll> d(n);
  rep(i, n) {
    ll c; cin >> c;
    sum0 += c;
    d[i] = p*(q-i-1) - c;
  }
  sort(d.rbegin(), d.rend());

  ll ans = sum0, sum1 = 0;
  rep(i, n) {
    sum1 += d[i];
    ll res = sum0 + sum1 + p*(i+1)*(i+1);
    ans = max(ans, res);
  }
  cout << ans << '\n';
}
タイトルとURLをコピーしました