Submission #71503

#	Time	Username	Problem	Language	Result	Execution time	Memory
71503		thiago4532	Chase (CEOI17_chase)	C++17	100 / 100	813 ms	357840 KiB

chase

#include <bits/stdc++.h>
 
using namespace std;
const int maxn = 1e5 + 10, maxk = 1e2 + 10;
vector<int> grafo[maxn];
int n, k;
int64_t gain[maxn], x[maxn], dp[maxn][maxk][2], value[maxn][maxk], copia[maxn][maxk], gain_copy[maxn];

inline void add(int u, int v){
	for(int i=1;i<=k;i++){
		if(value[v][i] > dp[u][i][0]){
			dp[u][i][1] = dp[u][i][0];
			dp[u][i][0] = value[v][i];
		}else if(value[v][i] > dp[u][i][1])
			dp[u][i][1] = value[v][i];
	}
	gain[u] += x[v];
}

inline void del(int u, int v){
	for(int i=1;i<=k;i++){
		if(dp[u][i][0] == value[v][i]){
			dp[u][i][0] = dp[u][i][1];
		}
	}
	gain[u] -= x[v];
}

inline void make_copy(int u){
	for(int i=0;i<=k;i++)
		copia[u][i] = dp[u][i][0];
	gain_copy[u] = gain[u];
}

inline void roll_back(int u){
	for(int i=0;i<=k;i++)
		dp[u][i][0] = copia[u][i];
	gain[u] = gain_copy[u];
}

inline void compute(int u){	
	for(int i=1;i<=k;i++)
		value[u][i] = max(dp[u][i][0], dp[u][i-1][0] + gain[u]);
}

void dfs(int u, int p=0){
	for(auto v : grafo[u]){
		if(v == p) continue; 
		dfs(v, u);
		add(u, v);
	}
	compute(u);
}

int64_t ans;
void rotate(int u, int p=0){
	compute(u);
	make_copy(u);
	ans = max(ans, value[u][k]);

	for(auto v : grafo[u]){
		if(v == p) continue;
		del(u, v);
		compute(u);
		add(v, u);
		rotate(v, u);
		roll_back(u);
	}
}

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(1);
	rotate(1);
	cout << ans << "\n";
	return 0;
}

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