Submission #105371

#	Time	Username	Problem	Language	Result	Execution time	Memory
105371		ryansee	Chase (CEOI17_chase)	C++14	100 / 100	2811 ms	335888 KiB

chase

#include "bits/stdc++.h"
using namespace std;
#define FAST ios::sync_with_stdio(false); cin.tie(0);
#define ll long long int
#define ld long double
#define pb push_back
#define ins insert
#define ph push
#define mmst(x,v) memset(x,v,sizeof(x))
#define all(x) (x).begin(), (x).end()
#define FOR(i,s,e) for(ll i=s;i<e;i++)
#define cbr cerr << "hi\n"
#define MAXN (100006)
#define LLINF ((long long)1e18)
typedef pair <ll, ll> pi;
typedef pair <ll, pi> spi;
int n, A[MAXN];
ll dp[MAXN][102], V,ans, alu[MAXN][102];
vector<int>v[MAXN];
// 0 - went down (aka. kid pressed) 
// 1 - ready to go up (aka. pressed)
void up(ll &x, ll y) { x=max(x,y); }
void dfs(ll x, ll p) {
    FOR(i,0,V+1) dp[x][i]=-LLINF,alu[x][i]=0;
	ll sum = 0; for(auto i:v[x])if(i!=p)sum+=A[i];
	//~ if(p == -1) {
		//~ dp[x][1][1] = sum; if(V>=1)ans = max(ans,sum);
		//~ dp[x][0][1] = 0;
	//~ } else {
		//~ FOR(i,0,V+1) {
			//~ up(dp[x][i][0],dp[p][i][1]);
			//~ up(dp[x][i][1],dp[p][i][0]); // np
			//~ up(dp[x][i][1],dp[p][i][1]);
			//~ if(i>0)up(dp[x][i][1],dp[p][i-1][0]+sum); // press
			//~ if(i>0)up(dp[x][i][1],dp[p][i-1][1]+sum);
//~ // 			up(dp[x][i][2],max({dp[p][i][2],dp[p][i][0]}));
			//~ ans=max({ans,dp[x][i][0],dp[x][i][1]});
		//~ }
	//~ }
	if(v[x].size()==1&&p!=-1) { dp[x][1] = sum;  dp[x][0] = 0; if(V>=1)ans=max(ans,sum);}
	alu[x][1] = sum+A[p]; if(V>=1)ans = max(ans,sum+A[p]);
	priority_queue <ll> pq[102];
	for(auto i : v[x]) if(i!=p) {
		dfs(i,x); // sum -= A[i];
		FOR(b,0,V+1) {
			up(dp[x][b], dp[i][b]); pq[b].ph(max(dp[i][b],dp[i][b]));
			up(dp[x][b], max(dp[i][b],dp[i][b])); 
			if(b>0)up(dp[x][b],max(dp[i][b-1],dp[i][b-1])+sum);
			ans=max({ans,dp[x][b],dp[x][b]});
		}
		//~ sum += A[i];
	}
	// if(x==6) cerr << x << ' ' << dp[x][1][0] << '\n';
	//~ 0 - went down
	//~ 1 - stayed
	if(p!=-1)sum += A[p];
	for(auto i : v[x]) if(i!=p) {
		sum -= A[i];
		FOR(b,0,V+1) {
			up(alu[x][b],alu[i][b]);
			up(alu[x][b],max(alu[i][b],alu[i][b]));
			if(b>0)up(alu[x][b],max(alu[i][b-1],alu[i][b-1])+sum);
			ans=max({ans,alu[x][b],alu[x][b]});
		}
		sum += A[i];
	}
	for(ll i = 1; i <= V; i++) dp[x][i]=max(dp[x][i],dp[x][i-1]), dp[x][i]=max(dp[x][i],dp[x][i-1]);
	for(ll i = 1; i <= V; i++) alu[x][i]=max(alu[x][i],alu[x][i-1]), alu[x][i]=max(alu[x][i],alu[x][i-1]);
	// if(x==6||x==7) cerr << x << ' ' << alu[x][1][0]
	FOR(b,0,V+1)pq[b].ph(-LLINF);
	for(auto i : v[x]) if(i!=p) { // will be my alu
		sum -= A[i];
		FOR(b,0,V+1) {
			bool zeroneed = 0, oneneed = 0; // ll hisdp = max({max(dp[i][b][0],dp[i][b][1])});
			if(pq[V-b].top() == max(dp[i][V-b],dp[i][V-b])) { pq[V-b].pop(); zeroneed = 1; }
			ll z0 = alu[i][b];
			ll o1 = max(alu[i][b],alu[i][b]);
			if(b>0)o1=max(o1,max(alu[i][b-1],alu[i][b-1])+sum+(b-1==0?A[i]:0ll));
			ans=max({ans,z0+pq[V-b].top(),o1+pq[V-b].top()});
			if(zeroneed) pq[V-b].ph(max(dp[i][V-b],dp[i][V-b]));
		}
		sum += A[i];
	}
	return;
}
 
int main()
{
	// freopen("in","r",stdin);
	FAST
	cin>>n>>V; if(n==1) { cout << 0 << '\n'; return 0; }
	FOR(i,1,n+1)cin>>A[i]; for(ll i=0;i<n-1;i++) { ll a, b; cin >> a >> b; v[a].pb(b); v[b].pb(a); }
// 	ll ans = 0;
	FOR(i,1,n+1) {
	    if(i>1)break;
// 		for(ll i=0;i<MAXN;i++) for(ll j=0;j<102;j++) dp[i][j][0]=dp[i][j][1]=dp[i][j][2]=-LLINF;
		dfs(i,-1);
	// ll ans = 0;
// 		for(ll i=1;i<=n;i++) for(ll j=0;j<=V;j++)ans=max({ans,dp[i][j][0],dp[i][j][1],dp[i][j][2]}); 
	}
	//~ dfs(6,-1);
	cout<<ans<<'\n';
}
/*
 * 
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
* 
*/

Compilation message (stderr)

chase.cpp: In function 'void dfs(long long int, long long int)':
chase.cpp:74:23: warning: unused variable 'oneneed' [-Wunused-variable]
    bool zeroneed = 0, oneneed = 0; // ll hisdp = max({max(dp[i][b][0],dp[i][b][1])});
                       ^~~~~~~
chase.cpp: In function 'int main()':
chase.cpp:11:20: warning: this 'for' clause does not guard... [-Wmisleading-indentation]
 #define FOR(i,s,e) for(ll i=s;i<e;i++)
                    ^
chase.cpp:92:2: note: in expansion of macro 'FOR'
  FOR(i,1,n+1)cin>>A[i]; for(ll i=0;i<n-1;i++) { ll a, b; cin >> a >> b; v[a].pb(b); v[b].pb(a); }
  ^~~
chase.cpp:92:25: note: ...this statement, but the latter is misleadingly indented as if it were guarded by the 'for'
  FOR(i,1,n+1)cin>>A[i]; for(ll i=0;i<n-1;i++) { ll a, b; cin >> a >> b; v[a].pb(b); v[b].pb(a); }
                         ^~~

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