Submission #59723

#	Time	Username	Problem	Language	Result	Execution time	Memory
59723		thiago4532	Chase (CEOI17_chase)	C++17	0 / 100	4108 ms	309492 KiB

chase

#include <bits/stdc++.h>
#define int long long

using namespace std;
const int maxn = 1e5 + 10, maxk = 1e2 + 10;
vector<int> grafo[maxn];
int n, k;
int gain[maxn], x[maxn], dp[maxn][maxk], dp2[maxn][maxk], ans[maxn];

void dfs(int u, int p=0){
	gain[u] = 0;
	for(auto v : grafo[u]){
		if(v == p) continue;
		gain[u] += x[v];

		dfs(v, u);
	}
}

void solve(int u, int p=0){
	for(int l=0;l<=k;l++)
		dp[u][l] = dp2[u][l] = 0;

	for(auto v : grafo[u]){
		if(v == p) continue;
		solve(v, u);

		for(int l=1;l<=k;l++){
			int x = max(dp[v][l], dp[v][l-1] + gain[u]);

			if(dp[u][l] < x) dp2[u][l] = dp[u][l], dp[u][l] = x;
			else if(dp2[u][l] < x) dp2[u][l] = x;
		}
	}
}


void rotate(int u, int p=0){
	ans[u] = dp[u][k];
	// cout << "ans(" << u << "): " << ans[u] << " " << dp2[u][k] << "\n";
	// cout << "Raiz em " << u << " com valor de dp " << dp[u][k] << " - " << dp2[u][k] << "\n";

	for(auto v : grafo[u]){
		if(v == p) continue;

		int x1 = gain[u], x2 = gain[v];

		int x3[k+1], x4[k+1], x5[k+1];
		for(int l=1;l<=k;l++) x3[l] = dp[u][l], x4[l] = dp[v][l], x5[l] = dp2[v][l];

		for(int l=1;l<=k;l++){
			if(dp[u][l] == dp[v][l] || dp[u][l] == dp[v][l-1] + gain[u])
				dp[u][l] = dp2[u][l];
			if(dp[u][l] != 0) dp[u][l] -= x[v];
		}

		gain[u] -= x[v];
		gain[v] += x[u];

		for(int l=0;l<=k;l++)
			dp[v][l] = dp2[v][l] = 0;

		for(auto b : grafo[v]){
			for(int l=1;l<=k;l++){
				if(dp[v][l] < dp[b][l]) dp2[v][l] = dp[v][l], dp[v][l] = dp[b][l];
				else if(dp2[v][l] < dp[b][l]) dp2[v][l] = dp[b][l];

				if(dp[v][l] < dp[b][l-1] + gain[v]) dp2[v][l] = dp[v][l], dp[v][l] = dp[b][l-1] + gain[v];
				else if(dp2[v][l] < dp[b][l-1] + gain[v]) dp2[v][l] = dp[b][l-1] + gain[v];
			}
		}

		rotate(v, u);
		gain[u] = x1, gain[v] = x2;
		for(int l=1;l<=k;l++) dp[u][l] = x3[l], dp[v][l] = x4[l], dp2[v][l] = x5[l];

	}
}

int32_t main(){
	ios::sync_with_stdio(false), cin.tie(0);
	cin >> n >> k;
	for(int i=1;i<=n;i++)
		cin >> x[i];

	for(int i=1;i<n;i++){
		int a, b;
		cin >> a >> b;
		grafo[a].push_back(b);
		grafo[b].push_back(a);
	}

	dfs(3);
	solve(3);

	rotate(3);

	int resp=0;
	for(int i=1;i<=n;i++){
		// cout << "gain: " << gain[i] << "\n";
		// cout << "dp: " << dp[i][k] << "\n";
		resp = max(resp, ans[i]);
	}

	int resp2=0;
	for(int i=1;i<=n;i++){
		dfs(i);
		solve(i);

		resp2 = max(resp2, dp[i][k]);
		 // cout << "ans2(" << i << "): " << dp[i][k] << " " << dp2[i][k] << "\n";
	}

	cout << resp << "\n";
	cout << resp2 << "\n";
	return 0;
}

• Subtask 1: 0 / 20
• Subtask 2: 0 / 20
• Subtask 3: 0 / 30
• Subtask 4: 0 / 30