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

原题链接：https://codeforces.ml/problemset/problem/711/C

题意

有n棵树，m种颜色，我们定义美丽度（如果相邻颜色不同则美丽数加一），cost[i][j]为给第i棵树上j种颜色的花费，问最后满足美丽度为k时的最小花费。

分析

根据状态很容易想到三维的DP，$dp[i][j][o]$代表前i棵树，当前美丽度为j，当前使用颜色为o时的最小花费，状态转移也可以轻松推出，复杂度是O(nkm^2)。但本题有几个坑点，就是如果首尾的颜色已经固定了，那么就不需要去遍历颜色了。

Code

#include <bits/stdc++.h>
using namespace std;
#define fi first
#define se second
#define re register
typedef long long ll;
typedef pair<ll, ll> PII;
typedef unsigned long long ull;
const int N = 1e6 + 20, M = 1e6 + 5, INF = 0x3f3f3f3f;
const int MOD = 1e9+9;
int a[N];
ll dp[105][105][105];
ll cost[105][105];
void solve() {
    int n, m, k; cin >> n >> m >> k;
    for (int i = 1; i <= n; i++) cin >> a[i];
    for (int i = 1; i <= n; i++) {
        for (int j = 1; j <= m; j++) {
            cin >> cost[i][j];
        }
    }
    memset(dp, 0x3f, sizeof dp);
    if (a[1]) {
        dp[1][1][a[1]] = 0;
    } else {
        for (int i = 1; i <= m; i++) dp[1][1][i] = cost[1][i];
    }

    for (int i = 2; i <= n; i++) {
        for (int j = 1; j <= i; j++) {
            if (a[i]) {
                for (int o = 1; o <= m; o++) {
                    if (a[i] == o) dp[i][j][o] = min(dp[i][j][o], dp[i - 1][j][o]);
                    else dp[i][j][a[i]] = min(dp[i][j][a[i]], dp[i - 1][j - 1][o]);
                }
            } else {
                for (int o = 1; o <= m; o++) {
                    for (int l = 1; l <= m; l++) {
                        if (o == l) dp[i][j][l] = min(dp[i][j][l], dp[i-1][j][l] + cost[i][l]);
                        else dp[i][j][l] = min(dp[i][j][l], dp[i-1][j-1][o] + cost[i][l]);
                    }
                }
            }
        }
    }
    ll ans = 1e18;
    if (a[n]) ans = min(ans, dp[n][k][a[n]]);
    else {
        for (int i = 1; i <= m; i++) {
            ans = min(ans, dp[n][k][i]);
        }
    }
    if (ans == 1e18) cout << -1 << endl;
    else cout << ans << endl;
}

signed main() {
    ios::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);
#ifdef ACM_LOCAL
    freopen("input", "r", stdin);
    freopen("output", "w", stdout);
#endif
    solve();
}