Codeforces 633F The Chocolate Spree 树形dp

The Chocolate Spree

对拍拍了半天才知道哪里写错了。。

dp[ i ][ j ][ k ]表示在 i 这棵子树中有 j 条链， 是否有链延伸上来。

#include<bits/stdc++.h>
#define LL long long
#define fi first
#define se second
#define mk make_pair
#define PLL pair<LL, LL>
#define PLI pair<LL, int>
#define PII pair<int, int>
#define SZ(x) ((int)x.size())
#define ull unsigned long long

using namespace std;

const int N = 1e5 + 7;
const int inf = 0x3f3f3f3f;
const LL INF = 0x3f3f3f3f3f3f3f3f;
const int mod = 998244353;
const double eps = 1e-6;
const double PI = acos(-1);

int n, a[N];
LL dp[N][3][2];

vector<int> G[N];

inline bool chkmax(LL &a, LL b) {
    return a < b ? a = b, true : false;
}

void dfs(int u, int fa) {
    LL tmp[5][3];
    LL gg[5][3];
    dp[u][0][0] = 0;
    for(int i = 0; i < 5; i++)
        for(int j = 0; j < 3; j++)
            tmp[i][j] = -INF;
    tmp[0][0] = 0;
    for(auto& v : G[u]) {
        if(v == fa) continue;
        dfs(v, u);
        memcpy(gg, tmp, sizeof(gg));
        for(int i = 0; i <= 2; i++) {
            for(int j = 0; j <= 2; j++) {
                for(int x = 0; x <= 2; x++) {
                    for(int y = 0; y <= 1; y++) {
                        if(!i && j || !x && y) continue;
                        if(j + y > 2) continue;
                        tmp[i + x][j + y] = max(tmp[i + x][j + y], gg[i][j] + dp[v][x][y]);
                    }
                }
            }
        }
    }
    // dp[1][0]
    chkmax(dp[u][1][0], tmp[0][0] + a[u]);
    chkmax(dp[u][1][0], tmp[1][0]);
    chkmax(dp[u][1][0], tmp[1][1] + a[u]);
    chkmax(dp[u][1][0], tmp[2][2] + a[u]);

    // dp[1][1]
    chkmax(dp[u][1][1], tmp[0][0] + a[u]);
    chkmax(dp[u][1][1], tmp[1][1] + a[u]);

    //dp[2][0]
    chkmax(dp[u][2][0], tmp[2][0]);
    chkmax(dp[u][2][0], tmp[2][1] + a[u]);
    chkmax(dp[u][2][0], tmp[3][2] + a[u]);

    //dp[2][1]
    chkmax(dp[u][2][1], tmp[2][1] + a[u]);
    chkmax(dp[u][2][1], tmp[2][2] + a[u]);
    chkmax(dp[u][2][1], tmp[1][0] + a[u]);
}

int main() {
//    freopen("test.in", "r", stdin);
    scanf("%d", &n);
    for(int i = 1; i <= n; i++) G[i].clear();
    for(int i = 1; i <= n; i++) scanf("%d", &a[i]);
    for(int i = 1; i < n; i++) {
        int u, v; scanf("%d%d", &u, &v);
        G[u].push_back(v);
        G[v].push_back(u);
    }
    for(int i = 1; i <= n; i++)
        for(int j = 0; j < 3; j++)
            for(int k = 0; k < 2; k++)
                dp[i][j][k] = -INF;
    dfs(1, 0);
    LL ans = 0;
    for(int i = 0; i <= 2; i++)
        for(int j = 0; j <= 1; j++)
            ans = max(ans, dp[1][i][j]);
    printf("%lld
", ans);
    return 0;
}

/*
*/