Submission #67978356

---

Source Code Expand

#include <bits/stdc++.h>
//#define int long long 
#define inf 0x3f3f3f3f
#define INF 0x3f3f3f3f3f3f3f3f
#define fr first
#define se second
#define endl '\n'
#define pb push_back
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int,int> pii;

const int N = 2e5 + 10;

int n, k;
ll a[N];
vector<int> G[N];

void chmax(ll& x, ll y){
    x = max(x, y);
}

void solve()
{
    cin >> n >> k;
    for(int i = 1; i <= n; i ++){
        cin >> a[i];
    }
    for(int i = 1; i <= n - 1; i ++){
        int u, v;
        cin >> u >> v;
        G[u].pb(v); G[v].pb(u);
    }

    vector<vector<vector<ll>>> dp(n + 1, vector<vector<ll>>(7, vector<ll>(2, -INF)));
    // dp[u][j][0]: 考虑以u为根的子树，不算点u，选取恰好j条互不重叠的路径，权值总和的最大值
    // dp[u][j][1]: 考虑以u为根的子树，算点u，选取恰好j条互不重叠的路径，权值总和的最大值
    vector<vector<ll>> f(n + 1, vector<ll>(7, -INF));
    // f[u][j]：考虑以u为根的子树，选取恰好j条互不重叠的路径，其中必须存在一条路径以u为端点，权值总和的最大值

    function<void(int,int)> DFS = [&](int u, int father){
        // init
        dp[u][0][0] = 0;
        dp[u][0][1] = 0;
        dp[u][1][1] = a[u];
        f[u][1] = a[u];
        for(auto v : G[u]){
            if(v == father) continue;
            DFS(v, u);
            // ！注意，在进行树形dp时，若想使用合并子树之前的结果转移，则最好（bixu）先将转移之前的结果保存下来，以避免在转移过程中被覆盖
            auto pre_dp_u = dp[u]; 
            auto pre_f_u = f[u];
            for(int j = 1; j <= k; j ++){
                for(int j1 = 0; j1 <= j; j1 ++){
                    int j2 = j - j1;
                    chmax(dp[u][j][1], pre_dp_u[j1][1] + dp[v][j2][1]);
                    chmax(dp[u][j][1], pre_f_u[j1] + f[v][j2 + 1]);
                    chmax(f[u][j], pre_f_u[j1] + dp[v][j2][1]);
                    chmax(f[u][j], (f[v][j1] + a[u]) + pre_dp_u[j2][0]);
                    chmax(dp[u][j][0], pre_dp_u[j1][0] + dp[v][j2][1]);
                }
            }
        }
        // for(int j = 0; j <= k; j ++){
        //     printf("dp[%d][%d][%d] = %lld\n", u, j, 0, dp[u][j][0]);
        //     printf("dp[%d][%d][%d] = %lld\n", u, j, 1, dp[u][j][1]);
        // }
    };
    DFS(1, 0);

    ll ans = 0;
    for(int i = 0; i <= k; i ++){
        ans = max(ans, dp[1][i][1]);
    }
    cout << ans << endl;
}


signed main()
{
    ios::sync_with_stdio(false); cin.tie(nullptr); cout.tie(nullptr);
    // int T=1; cin>>T; while(T--)
    solve();
    return 0;
}

Submission Info

Judge Result