CF 711C Coloring Trees(三维dp)

原题地址: https://codeforces.com/problemset/problem/711/C

题意:有$n$棵树,$m$种颜色，让你分成$k$段，数字$0$代表该树没没有进行染色，需要你从$m$种颜色中选择一个为其染色，染色的过程中需要消耗代价，求最小代价（已有颜色不必再染色）

思路:三维$dp$,定义状态$dp[i][j][k]$表示前$i$个树,分成$j$段,最后颜色是$k$的方案数。

需要注意的是初始化，对第一棵树是否有颜色分类讨论即可。

#include <bits/stdc++.h>
#define eps 1e-8
#define PI acos(-1)
#define lson l,mid,rt<<1
#define rson mid+1,r,(rt<<1)+1
#define CLR(x,y) memset((x),y,sizeof(x))
#define fuck(x) cerr << #x << "=" << x << endl
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
const int seed = 131;
const ll INF = 0x3f3f3f3f3f3f;
const int maxn = 1e5 + 5;
const int mod = 1e9 + 7;
int col[maxn];
int a[105][105];
ll dp[105][104][105];//dp[i][j][k]表示前i个,分成j段,最后颜色是k的方案数
int n, m, K;
int main() {
    scanf("%d%d%d", &n, &m, &K);//树,颜色数,段
    for (int i = 1; i <= n; i++) {
        scanf("%d", &col[i]);
    }
    for (int i = 1; i <= n; i++) {
        for (int j = 1; j <= m; j++) {
            scanf("%d", &a[i][j]);
        }
    }
    CLR(dp, INF);
    fuck(dp[1][1][1]);
    if (col[1]) {
        dp[1][1][col[1]] = 0;
    } else {
        for (int i = 1; i <= m; i++) {
            dp[1][1][i] = a[1][i];
        }
    }
    for (int i = 2; i <= n; i++) {
        for (int j = 1; j <= K; j++) {
            if (col[i]) {//有颜色
                for (int k = 1; k <= m; k++) {
                    if (col[i] == k) {
                        dp[i][j][col[i]] = min(dp[i][j][col[i]], dp[i - 1][j][col[i]]);
                    } else {
                        dp[i][j][col[i]] = min(dp[i][j][col[i]], dp[i - 1][j - 1][k]);
                    }
                }
            } else {//没有颜色
                for (int k = 1; k <= m; k++) {
                    for (int p = 1; p <= m; p++) {
                        if (k == p)dp[i][j][k] = min(dp[i][j][k], dp[i - 1][j][k] + a[i][k]);
                        else dp[i][j][k] = min(dp[i][j][k], dp[i - 1][j - 1][p] + a[i][k]);
                    }
                }
            }
        }
    }
    ll ans = INF;
    for (int i = 1; i <= m; i++) {
        ans = min(ans, dp[n][K][i]);
    }
    if (ans == INF) printf("-1\n");
    else printf("%lld\n", ans);
    return 0;
}