Submission #163342

#	Time	Username	Problem	Language	Result	Execution time	Memory
163342		alishahali1382	Chase (CEOI17_chase)	C++14	100 / 100	558 ms	135800 KiB

chase

#include <bits/stdc++.h>
#pragma GCC optimize ("O2")
#pragma GCC optimize ("unroll-loops")
//#pragma GCC optimize("no-stack-protector,fast-math")

using namespace std;
typedef long long ll;
typedef long double ld;
typedef pair<int, int> pii;
typedef pair<pii, int> piii;
typedef pair<ll, ll> pll;
#define debug(x) cerr<<#x<<'='<<(x)<<endl;
#define debugp(x) cerr<<#x<<"= {"<<(x.first)<<", "<<(x.second)<<"}"<<endl;
#define debug2(x, y) cerr<<"{"<<#x<<", "<<#y<<"} = {"<<(x)<<", "<<(y)<<"}"<<endl;
#define debugv(v) cerr<<#v<<" : ";for (auto x:v) cerr<<x<<' ';cerr<<endl;
#define all(x) x.begin(), x.end()
#define pb push_back
#define kill(x) return cout<<x<<'\n', 0;

const ld eps=1e-7;
const int inf=1000000010;
const ll INF=10000000000000010LL;
const int mod = 1000000007;
const int MAXN = 100010, K=101;

ll n, m, k, u, v, x, y, t, a, b, ans;
ll A[MAXN];
ll val[MAXN];
ll dp[MAXN][K];
int ind[MAXN][K][2];
vector<int> G[MAXN];

void calcdp(int node){
	for (int j=1; j<=k; j++) dp[node][j]=max(dp[ind[node][j][0]][j], dp[ind[node][j-1][0]][j-1] + val[node]);
}

ll dfs1(int node, int par){
	for (int v:G[node]) if (v!=par) val[node]+=dfs1(v, node);
	
	for (int v:G[node]) if (v!=par) for (int j=1; j<=k; j++) if (dp[v][j]>dp[ind[node][j][0]][j]) ind[node][j][0]=v;
	for (int v:G[node]) if (v!=par) for (int j=1; j<=k; j++) if (v!=ind[node][j][0] && dp[v][j]>dp[ind[node][j][1]][j]) ind[node][j][1]=v;
	
	calcdp(node);
	return A[node];
}

void dfs2(int node, int par){	
	calcdp(node);
	for (int j=1; j<=k; j++) if (par){
		dp[node][j]=max(dp[node][j], dp[par][j]);
		dp[node][j]=max(dp[node][j], dp[par][j-1] + val[node]);
	}
	
	for (int j=1; j<=k; j++) ans=max(ans, dp[node][j]);
	
	for (int v:G[node]) if (v!=par){
		val[node]-=A[v];
		val[v]+=A[node];
		
		for (int j=1; j<=k; j++){
			if (ind[node][j][0]!=v) dp[node][j]=dp[ind[node][j][0]][j];
			else dp[node][j]=dp[ind[node][j][1]][j];
			
			if (ind[node][j-1][0]!=v) dp[node][j]=max(dp[node][j], dp[ind[node][j-1][0]][j-1] + val[node]);
			else dp[node][j]=max(dp[node][j], dp[ind[node][j-1][1]][j-1] + val[node]);
			
			if (par){
				dp[node][j]=max(dp[node][j], dp[par][j]);
				dp[node][j]=max(dp[node][j], dp[par][j-1] + val[node]);
			}
		}
		
		dfs2(v, node);
		
		val[v]-=A[node];
		val[node]+=A[v];
		calcdp(v);
		calcdp(node);
		for (int j=1; j<=k; j++) if (par){
			dp[node][j]=max(dp[node][j], dp[par][j]);
			dp[node][j]=max(dp[node][j], dp[par][j-1] + val[node]);
		}
	}
}

int main(){
	ios_base::sync_with_stdio(false);cin.tie(0);cout.tie(0);
	//freopen("input.txt", "r", stdin);
	//freopen("output.txt", "w", stdout);
	cin>>n>>k;
	for (int i=1; i<=n; i++) cin>>A[i];
	for (int i=1; i<n; i++){
		cin>>u>>v;
		G[u].pb(v);
		G[v].pb(u);
	}
	dfs1(1, 0);
	dfs2(1, 0);
	cout<<ans<<'\n';
	
	return 0;
}
/*
12 2
2 3 3 8 1 5 6 7 8 3 5 4
2 1
2 7
3 4
4 7
7 6
5 6
6 8
6 9
7 10
10 11
10 12

*/

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