This submission is migrated from previous version of oj.uz, which used different machine for grading. This submission may have different result if resubmitted.
#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;
ll n, A[MAXN], dp[MAXN][102][3], V,ans, alu[MAXN][102][3];
vector<ll>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][0]=dp[x][i][1]=-LLINF,alu[x][i][0]=alu[x][i][1]=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][1] = sum; dp[x][0][1] = 0; if(V>=1)ans=max(ans,sum);}
alu[x][1][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][0], dp[i][b][1]); pq[b].ph(max(dp[i][b][1],dp[i][b][0]));
up(dp[x][b][1], max(dp[i][b][0],dp[i][b][1]));
if(b>0)up(dp[x][b][1],max(dp[i][b-1][0],dp[i][b-1][1])+sum);
ans=max({ans,dp[x][b][0],dp[x][b][1]});
}
//~ 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][0],alu[i][b][1]);
up(alu[x][b][1],max(alu[i][b][0],alu[i][b][1]));
if(b>0)up(alu[x][b][1],max(alu[i][b-1][0],alu[i][b-1][1])+sum);
ans=max({ans,alu[x][b][0],alu[x][b][1]});
}
sum += A[i];
}
for(ll i = 1; i <= V; i++) dp[x][i][0]=max(dp[x][i][0],dp[x][i-1][0]), dp[x][i][1]=max(dp[x][i][1],dp[x][i-1][1]);
for(ll i = 1; i <= V; i++) alu[x][i][0]=max(alu[x][i][0],alu[x][i-1][0]), alu[x][i][1]=max(alu[x][i][1],alu[x][i-1][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][0],dp[i][V-b][1])) { pq[V-b].pop(); zeroneed = 1; }
if(b < V && pq[V-b-1].top() == max(dp[i][V-b-1][0],dp[i][V-b-1][1])) { pq[V-b-1].pop(); oneneed = 1; }
ll z0 = alu[i][b][1];
ll o1 = max(alu[i][b][0],alu[i][b][1]);
if(b>0)o1=max(o1,max(alu[i][b-1][0],alu[i][b-1][1])+sum);
ans=max({ans,z0+pq[V-b].top(),o1+(b<V?pq[V-b-1].top():-LLINF)}); // o1 needs to V-b-1 as u need to use one at x
if(zeroneed) pq[V-b].ph(max(dp[i][V-b][0],dp[i][V-b][1]));
if(oneneed) pq[V-b-1].ph(max(dp[i][V-b-1][0],dp[i][V-b-1][1]));
}
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); // cerr << alu[3][1][1] << '\n';
// 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 '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:93: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:93: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); }
^~~
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |