Codeforces Round #369 (Div. 2) C. Coloring Trees

問題

codeforces.com

考察

隣接している木をなるべく違う色で塗る必要があります。
木の数、色の数は100と小さいので、各色の塗りを試しても間に合いそうです。
一方、必要なペイント量は大きな値です。
そこで、以下のようなDPを考えます。

dp[i][j][k] := i番目の木をjで塗ってi番目までの美しさがkのとき、必要なペイントの最小値

もともと塗られている木では、その色でのみ遷移を行うことで上手く処理ができます。

このようなDPでは、 O(n^{2} * m)のメモリが必要です。
漸化式の遷移では各色を試すので、全体の計算量は O(n^{2} * m^{2})です。
少し怪しい感じの計算量ですが、C++なら通りました。

コード

int main() {
    int n, m, k;
    cin >> n >> m >> k;

    vi def(n);
    rep(i, n)cin >> def[i];

    vvi amt(n, vi(m+1));
    rep(i, n) {
        rep(j, m) {
            cin >> amt[i][j];
        }
    }

    const ll INF = 1ll << 60;

    vector<vvll> dp(n + 1, vvll(m + 1, vll(n + 2, INF)));
    dp[0][0][0] = 0;

    rep(i, n) {
        rep(j, m + 1) {
            rep(k, n ) {
                if (dp[i][j][k] == INF)
                    continue;

                if (def[i] != 0) {
                    int t = k;
                    if (j != def[i])
                        t++;
                    chmin(dp[i + 1][def[i]][t], dp[i][j][k]);
                } else {
                    REP(c, 1, m + 1) {
                        int t = k;
                        if (j != c)
                            t++;
                        chmin(dp[i + 1][c][t], dp[i][j][k] + amt[i][c - 1]);
                    }
                }
            }
        }
    }

    ll ans = INF;
    REP(i, 1, m+1)
        chmin(ans, dp[n][i][k]);

    if(ans == INF)
        cout << -1 << endl;
    else
        cout << ans << endl;


    return 0;
}