Codeforces Round #595 (Div. 3) F. Maximum Weight Subset 树形dp

原题链接：https://codeforces.com/problemset/problem/1249/F

题意

有一颗树，每个节点都有权值，在树上选择距离大于k且选取集合权值和最大是多少。

分析

因为要选择大于k距离的点对，因此状态必须记录在子节点中距离大于j的最大权值，那么状态就可以直接出来了。我们设$f[i][j]$代表以i为根的子树中，距离大于j的最大权值。

分类讨论一下

1. $j==0时，f[i][0]=a[i]+dp[v][k],v\in i的子节点$
2. $j>0时，f[i][j]=a[i]+max(dp[v1][max(j-1,k-j)]+dp[v2][j-1])j\in [1,n)$

最后不要忘记从后往前取一遍max，因为我们的状态为大于j的最大权值，也就是说前面的状态是包含后面的状态的，因此后面的状态不要忘了更新前面的状态。

Code

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int, int> PII;
const ll INF = 1e18;
const int N = 4e5 + 10;
const int M = 1000010;
const int MOD = 1e9 + 7;
int a[205];
int n, k;
ll dp[205][205];
vector<int> g[205];
void dfs(int x, int fa) {
    dp[x][0] = a[x];
    for (auto v : g[x]) {
        if (v == fa) continue;
        dfs(v, x);
        dp[x][0] += dp[v][k];
    }
    for (int i = 1; i < n; i++) {
        for (auto v1 : g[x]) {
            if (v1 == fa) continue;
            ll tmp = dp[v1][i-1];
            for (auto v2 : g[x]) {
                if (v2 == fa || v2 == v1) continue;
                tmp += dp[v2][max(i-1, k-i)];
            }
            dp[x][i] = max(dp[x][i], tmp);
        }
    }
    for (int i = n-1; i >= 0; i--) dp[x][i] = max(dp[x][i], dp[x][i+1]);
}
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].push_back(v);
        g[v].push_back(u);
    }
    dfs(1, 0);
    cout << dp[1][0] << endl;
}
signed main() {
    ios_base::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
#ifdef ACM_LOCAL
    freopen("input", "r", stdin);
    freopen("output", "w", stdout);
    signed test_index_for_debug = 1;
    char acm_local_for_debug = 0;
    do {
        if (acm_local_for_debug == '$') exit(0);
        if (test_index_for_debug > 20)
            throw runtime_error("Check the stdin!!!");
        auto start_clock_for_debug = clock();
        solve();
        auto end_clock_for_debug = clock();
        cout << "Test " << test_index_for_debug << " successful" << endl;
        cerr << "Test " << test_index_for_debug++ << " Run Time: "
             << double(end_clock_for_debug - start_clock_for_debug) / CLOCKS_PER_SEC << "s" << endl;
        cout << "--------------------------------------------------" << endl;
    } while (cin >> acm_local_for_debug && cin.putback(acm_local_for_debug));
#else
    solve();
#endif
    return 0;
}