1249F - Maximum Weight Subset - CodeForces Solution

#include <bits/stdc++.h>
using namespace std;
//(>'-'<)
typedef long long ll;

const int N = 5e3 + 2;
const int oo = 1e9 + 7;

int n, k, a[N];
int sz[N], dp[N][N], tmp[N], suf_max_u[N], suf_max_v[N];
vector<int> adj[N];

bool maximize(int &a, int b) { return a < b ? a = b, true : false; }

void dfs(int u, int p) {
    sz[u] = 1;
    dp[u][0] = a[u];
    for(int &v : adj[u]) if(v != p) {
        dfs(v, u);
        for(int h = 0; h <= sz[u] + sz[v]; ++h) {
            if (h > 0) tmp[h] = max(dp[u][h], dp[v][h - 1]);
            else tmp[h] = dp[u][h];
        }
        suf_max_u[sz[u] + 1] = -oo;
        for(int h = sz[u]; h >= 0; --h) {
            suf_max_u[h] = max(dp[u][h], suf_max_u[h + 1]);
        }
        suf_max_v[sz[v] + 1] = -oo;
        for(int h = sz[v]; h >= 0; --h) {
            suf_max_v[h] = max(dp[v][h], suf_max_v[h + 1]);
        }
        for(int h1 = 0; h1 <= sz[u]; ++h1) {
            int h2 = max(k - h1, h1 - 1);
            if (h2 <= sz[v]) maximize(tmp[h1], dp[u][h1] + suf_max_v[h2]);
        }
        for(int h2 = 0; h2 <= sz[v]; ++h2) {
            int h1 = max(k - h2, h2 + 1);
            if (h1 <= sz[u]) maximize(tmp[h2 + 1], suf_max_u[h1] + dp[v][h2]);
        }
        sz[u] += sz[v];
        for(int h = 0; h <= sz[u]; ++h) {
            dp[u][h] = tmp[h];
            tmp[h] = -oo;
        }
    }
    // for(int i = 0; i <= sz[u]; ++i) {
    //     cout << u + 1 << ' ' << i << ' ' << dp[u][i] << '\n';
    // }
    // cout << '\n';
}

int main() {
    ios_base::sync_with_stdio(0);cin.tie(0);

    cin >> n >> k;
    for(int i = 0; i < n; ++i) cin >> a[i];
    for(int i = 1; i < n; ++i) {
        // int x; cin >> x;
        // adj[i].push_back(x);
        // adj[x].push_back(i);
        int u, v; cin >> u >> v;
        --u, --v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    memset(dp, -0x3f, sizeof dp);
    memset(tmp, -0x3f, sizeof tmp);
    dfs(0, -1);
    int res = 0;
    for(int h = 0; h <= n; ++h) {
        res = max(res, dp[0][h]);
    }
    cout << res;

    return 0;
}