#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
vector<vector<ll>> adjlist;
vector<ll> neighbours;
vector<ll> val;
vector<vector<pair<ll,ll>>> upv;
vector<vector<pair<ll,ll>>> downv;
ll v;
ll INF = 1e18;
ll down(ll node, ll par, ll c, ll k){
if (k<=0){
return 0;
}
ll s = -1;
if (c!=1){
s=(c==0)?0:2;
}
if (s!=-1){
if (s==0){
if (downv[node][k].first!=-1){
return downv[node][k].first;
}
} else {
if (downv[node][k].second!=-1){
return downv[node][k].second;
}
}
}
ll mx = 0;
if (c==0){
if (adjlist[node].size()==1&&adjlist[node][0]==par){
mx=max(mx,neighbours[node]+val[node]);
}
for (ll i : adjlist[node]){
if (i==par){continue;}
mx=max(
mx,
max(
down(i,node,0,k-1)-val[i],
max(
down(i,node,1,k-1)-val[i]+val[node],
down(i,node,2,k-1)+val[node]
)
)+neighbours[node]);
}
} else if (c==1){
for (ll i : adjlist[node]){
if (i==par){continue;}
mx=max(mx,down(i,node,0,k)-val[i]);
}
} else {
for (ll i : adjlist[node]){
if (i==par){continue;}
mx=max(mx,max(down(i,node,1,k),down(i,node,2,k)));
}
}
if (s!=-1){
if (s==0){
return downv[node][k].first=mx;
} else {
return downv[node][k].second=mx;
}
} else {
return mx;
}
}
ll up(ll node, ll par, ll c, ll k){
if (k<=0){
return 0;
}
ll s = -1;
if (c!=1){
s=(c==0)?0:2;
}
if (s!=-1){
if (s==0){
if (upv[node][k].first!=-1){
return upv[node][k].first;
}
} else {
if (upv[node][k].second!=-1){
return upv[node][k].second;
}
}
}
ll mx = 0;
if (c==0){
if (adjlist[node].size()==1&&adjlist[node][0]==par){
mx=max(mx,neighbours[node]);
}
for (ll i : adjlist[node]){
if (i==par){continue;}
mx=max(
mx,
max(
up(i,node,0,k-1)-val[i],
max(
up(i,node,1,k-1)-val[i],
up(i,node,2,k-1)
)
)+neighbours[node]);
}
} else if (c==1){
for (ll i : adjlist[node]){
if (i==par){continue;}
mx=max(mx,up(i,node,0,k));
}
} else {
for (ll i : adjlist[node]){
if (i==par){continue;}
mx=max(mx,max(up(i,node,1,k),up(i,node,2,k)));
}
}
//cout<<"up "<<node+1<<" "<<c<<" "<<k<<" "<<mx<<'\n';
if (s!=-1){
if (s==0){
return upv[node][k].first=mx;
} else {
return upv[node][k].second=mx;
}
} else {
return mx;
}
}
ll dp(ll node, ll par){
ll mx = max(
up(node,par,0,v),
max(up(node,par,1,v),up(node,par,2,v)));
mx=max(mx,
max(
down(node,par,0,v)-val[node],
max(down(node,par,1,v),down(node,par,2,v))
)
);
//cout<<"pre dp 1 "<<node+1<<" "<<mx<<'\n';
for (ll i : adjlist[node]){
if (i==par){continue;}
mx=max(mx,dp(i,node));
}
// middle not used
for (ll k = 1; k<v; k++){
ll maxu1=-INF, maxu1i=0, maxu2=-INF, maxu2i=0,
maxd1=-INF, maxd1i=0, maxd2=-INF, maxd2i=0;
for (ll i : adjlist[node]){
if (i==par){continue;}
ll uv = max(up(i,node,1,k),up(i,node,2,k));
ll dv = max(down(i,node,1,v-k),down(i,node,2,v-k));
if (uv>=maxu1){
maxu1=uv;maxu1i=i;
} else if (uv>=maxu2){
maxu2=uv;maxu2i=i;
}
if (dv>=maxu1){
maxd1=dv;maxd1i=i;
} else if (dv>=maxd2){
maxd2=dv;maxd2i=i;
}
}
if (maxd1i!=maxu1i){
mx=max(mx,maxu1+maxd1);
} else {
mx=max(mx,max(maxu1+maxd2,maxu2+maxd1));
}
}
//cout<<"pre dp 2 "<<node+1<<" "<<mx<<'\n';
//middle used
if (v>=3){
for (ll k = 1; k<v-1; k++){
//if centre is start
ll maxu1=-INF, maxu1i=0, maxu2=-INF, maxu2i=0,
maxd1=-INF, maxd1i=0, maxd2=-INF, maxd2i=0;
for (ll i : adjlist[node]){
if (i==par){continue;}
ll uv = max(up(i,node,1,k)-val[i],up(i,node,1,k));
ll dv = max(down(i,node,0,v-1-k)-val[node]-val[i],
max(down(i,node,1,v-1-k)-val[i],down(i,node,2,v-1-k)));
if (uv>=maxu1){
maxu1=uv;maxu1i=i;
} else if (uv>=maxu2){
maxu2=uv;maxu2i=i;
}
if (dv>=maxu1){
maxd1=dv;maxd1i=i;
} else if (dv>=maxd2){
maxd2=dv;maxd2i=i;
}
}
if (maxd1i!=maxu1i){
mx=max(mx,maxu1+maxd1+neighbours[node]);
} else {
mx=max(mx,max(maxu1+maxd2,maxu2+maxd1)+neighbours[node]);
}
}
//cout<<"pre dp 3 "<<node+1<<" "<<mx<<'\n';
for (ll k = 1; k<v-1; k++){
//if centre is middle
ll maxu1=-INF, maxu1i=0, maxu2=-INF, maxu2i=0,
maxd1=-INF, maxd1i=0, maxd2=-INF, maxd2i=0;
for (ll i : adjlist[node]){
if (i==par){continue;}
ll uv = up(i,node,0,k);
ll dv = max(down(i,node,0,v-1-k)-val[node]-val[i],
max(down(i,node,1,v-1-k)-val[i],down(i,node,2,v-1-k)));
if (uv>=maxu1){
maxu1=uv;maxu1i=i;
} else if (uv>=maxu2){
maxu2=uv;maxu2i=i;
}
if (dv>=maxu1){
maxd1=dv;maxd1i=i;
} else if (dv>=maxd2){
maxd2=dv;maxd2i=i;
}
if (maxd1i!=maxu1i){
mx=max(mx,maxu1+maxd1+neighbours[node]+val[node]);
} else {
mx=max(mx,max(maxu1+maxd2,maxu2+maxd1)+neighbours[node]+val[node]);
}
}
}
}
//cout<<"dp "<<node+1<<" "<<mx<<'\n';
return mx;
}
int main(){
cin.tie(0);
ios_base::sync_with_stdio(0);
ll n;
cin>>n>>v;
adjlist.resize(n);
neighbours.resize(n);
val.resize(n);
upv.resize(n,vector<pair<ll,ll>>(v+1,pair<ll,ll>({-1,-1})));
downv.resize(n,vector<pair<ll,ll>>(v+1,pair<ll,ll>({-1,-1})));
for (ll i = 0; i<n; i++){
cin>>val[i];
}
for (ll i = 0; i<n-1; i++){
ll a,b;
cin>>a>>b;
a--;b--;
adjlist[a].push_back(b);
adjlist[b].push_back(a);
neighbours[a]+=val[b];
neighbours[b]+=val[a];
}
cout<<dp(0,-1)<<'\n';
}
Compilation message
chase.cpp: In function 'll dp(ll, ll)':
chase.cpp:148:40: warning: variable 'maxu2i' set but not used [-Wunused-but-set-variable]
ll maxu1=-INF, maxu1i=0, maxu2=-INF, maxu2i=0,
^~~~~~
chase.cpp:149:37: warning: variable 'maxd2i' set but not used [-Wunused-but-set-variable]
maxd1=-INF, maxd1i=0, maxd2=-INF, maxd2i=0;
^~~~~~
chase.cpp:176:41: warning: variable 'maxu2i' set but not used [-Wunused-but-set-variable]
ll maxu1=-INF, maxu1i=0, maxu2=-INF, maxu2i=0,
^~~~~~
chase.cpp:177:38: warning: variable 'maxd2i' set but not used [-Wunused-but-set-variable]
maxd1=-INF, maxd1i=0, maxd2=-INF, maxd2i=0;
^~~~~~
chase.cpp:203:41: warning: variable 'maxu2i' set but not used [-Wunused-but-set-variable]
ll maxu1=-INF, maxu1i=0, maxu2=-INF, maxu2i=0,
^~~~~~
chase.cpp:204:38: warning: variable 'maxd2i' set but not used [-Wunused-but-set-variable]
maxd1=-INF, maxd1i=0, maxd2=-INF, maxd2i=0;
^~~~~~
| # |
결과 |
실행 시간 |
메모리 |
Grader output |
| 1 |
Incorrect |
5 ms |
384 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
결과 |
실행 시간 |
메모리 |
Grader output |
| 1 |
Incorrect |
5 ms |
384 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
결과 |
실행 시간 |
메모리 |
Grader output |
| 1 |
Incorrect |
3141 ms |
342244 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
결과 |
실행 시간 |
메모리 |
Grader output |
| 1 |
Incorrect |
5 ms |
384 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
