Facebook Hacker Cup 2017 Round 1 - Pie Progress

問題

Pie Progress | Facebook Hacker Cup 2017 Round 1

考察

購入する必要があるパイの個数は、日数と同じです。
N日間の中のどこでいくつ買うかを決定するので、以下のDPで処理できそうです。

dp[i][j] := i日目までにj個パイを買った場合の最小コスト

各日数でそれぞれのパイを買うかを試したいですが、普通にやっては計算量が増えすぎます。
そこで、「高いパイを買う必要はない」「追加コストは個数にのみ依存」という2点を活かすと、以下のような更新ができます。

  1. 最小コストのパイを取ってくる
  2. コストから00を引き、11を足す
  3. 1個買ったものとして、DP表を更新
  4. 次に小さいコストのパイを取ってくる
  5. コストから11を引き、22を足す
  6. 2個かったものとして、DP表を更新
  7. ...

このように、直前の個数の追加コストを引いてから新しく足すと、順番に処理できます。
更新し終わったDP表のdp[N][N]が答えになります。
priority_queueなどを使えば取り出しはO(logN)なので、全体で  O(T * N^{2} * logN) となり間に合いそうです。

なお、注意点として、「餓死するケース」を考える必要があります。
i日目の時点で買ったパイがi個未満だと餓死するので、そのケースは答えに含めないようにします。

コード

ll solve() {
    int N, M;
    cin >> N >> M;

    const int mN = N + 10;

    vector< priority_queue<int, vi, greater<int> > > pies(N);
    rep(i, N) {
        rep(j, M) {
            int t;
            cin >> t;
            pies[i].push(t);
        }
    }

    const ll INF = 1ll<<60;

    vvll dp(mN, vll(mN, INF));
    dp[0][0] = 0;

    rep(i, N) {
        // 0
        rep(j, mN)
            dp[i+1][j] = dp[i][j];

        const int sz = (int)pies[i].size();

        int cost = 0;
        rep(k, sz) {
            cost += pies[i].top();
            pies[i].pop();
            cost -= k*k;
            cost += (k+1)*(k+1);

            rep(j, mN) {
                int next = j + k + 1;
                if (dp[i][j] != INF && next < mN) {
                    chmin(dp[i+1][next], dp[i][j] + cost);
                }
            }
        }

        // check
        rep(j, i+1) {
            dp[i+1][j] = INF;
        }
    }

    return dp[N][N];
}

int main() {
    int T;
    cin >> T;
    rep(__t, T) {
        cout << "Case #" << __t + 1 << ": ";
        cout << solve() << endl;
    }

    return 0;
}

感想

思考停止DPで解きました。
オーダーがこれより大きくても、実行時間にゆとりはあるので通せそうです。