提出 #67966529

---

ソースコード 拡げる

#include <bits/stdc++.h>
using namespace std;

const long long inf = 1e18+3;

void solve() {
	int n, k;
	cin >> n >> k;
	vector<long long> a(n+1);
	for (int i=1; i<=n; i++) cin >> a[i];
	vector<vector<int>> adjl(n+1);
	for (int i=1; i<n; i++) {
		int u, v;
		cin >> u >> v;
		adjl[u].push_back(v);
		adjl[v].push_back(u);
	}

	vector<vector<vector<long long>>> dp(n+1, vector<vector<long long>>(k+1, vector<long long>(2, -inf)));
	auto dfs = [&](auto &self, int u, int p) -> void {
		dp[u][0][0] = dp[u][1][1] = a[u];
		dp[u][0][1] = 0;
		
		for (int &v : adjl[u]) {
			if (v == p) {
				swap(v, adjl[u][0]);
				break;
			}
		}

		int m = (int)adjl[u].size();
		int s = (u == 1 ? 0 : 1);

		for (int i=s; i<m; i++) {
			int v = adjl[u][i];
			self(self, v, u);
		}

		vector<vector<vector<long long>>> cdp(m+1, vector<vector<long long>>(k+1, vector<long long>(3, -inf)));
		cdp[s][0][0] = 0;

		for (int i=s+1; i<=m; i++) {
			int v = adjl[u][i-1];
			for (int ck=0; ck<=k; ck++) {
				// not merge path with i-th child
				for (int pk=0; pk<=ck; pk++) {
					for (int l=0; l<3; l++) {
						cdp[i][ck][l] = max(cdp[i][ck][l], cdp[i-1][pk][l] + dp[v][ck-pk][1]);
					}
				}
				for (int pk=0; pk<ck; pk++) {
					for (int l=0; l<3; l++) {
						cdp[i][ck][l] = max(cdp[i][ck][l], cdp[i-1][pk][l] + dp[v][ck-pk-1][0]);
					}
				}

				// merge path with i-th child
				for (int pk=0; pk<=ck; pk++) {
					for (int l=1; l<3; l++) {
						cdp[i][ck][l] = max(cdp[i][ck][l], cdp[i-1][pk][l-1] + dp[v][ck-pk][0]);
					}
				}
			}
		}

		for (int i=0; i<=k; i++) {
			dp[u][i][0] = max({dp[u][i][0], cdp[m][i][0] + a[u], cdp[m][i][1] + a[u]});
			dp[u][i][1] = max(dp[u][i][1], cdp[m][i][0]);
			if (i > 0) {
				dp[u][i][1] = max({dp[u][i][1], cdp[m][i-1][0] + a[u], cdp[m][i-1][1] + a[u], cdp[m][i-1][2] + a[u]});
			}
		}
	};
	dfs(dfs, 1, 0);

	long long ans = -inf;
	for (int i=0; i<=k; i++) {
		ans = max(ans, dp[1][i][1]);
		if (i > 0) ans = max(ans, dp[1][i-1][0]);
	}

	cout << ans << '\n';
}

int main() {
	ios_base::sync_with_stdio(0); cin.tie(0);
	int tcs = 1;
	// cin >> tcs;
	while (tcs--) {
		solve();
	}
	return 0;
}

提出情報

ジャッジ結果