# | Time | Username | Problem | Language | Result | Execution time | Memory |
---|---|---|---|---|---|---|---|
409272 | Hazem | Chase (CEOI17_chase) | C++14 | 0 ms | 0 KiB |
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 LL long long
#define F first
#define S second
#define pii pair<int,int>
#define piii pair<pair<int,int>,int>
const int N = 2e5+10;
const int M = 3e2+10;
const LL INF = 1e9;
const LL LINF = 2e18;
const LL MOD = 4294967296;
const double PI = 3.141592653589793;
LL a[N],dp[N][M],prevals[N][M];
pii mx[N][M][2],mxnode[N][M][2];
vector<int>adj[N];
int n,k;
void dfs(int i,int pre){
LL sum = 0;
for(auto x:adj[i])
if(x!=pre)
dfs(x,i),sum += a[x];
for(auto x:adj[i]){
if(x==pre)continue;
for(int j=1;j<=100;j++)
dp[i][j] = max(dp[i][j],max(dp[x][j],dp[x][j-1]+sum));
}
}
LL ans = 0;
void rootshift(int i,int pre){
LL sum = 0;
for(auto x:adj[i])
sum += a[x];
for(auto x:adj[i])
for(int j=1;j<=100;j++)
dp[i][j] = max(dp[i][j],max(dp[x][j],dp[x][j-1]+sum));
ans = max(ans,dp[i][k]);
for(auto x:adj[i]){
for(int j=1;j<=100;j++)
for(int j1=0;j1<2;j1++){
LL val = dp[x][j-j1];
if(val>mx[i][j][j1].F){
swap(mx[i][j][j1].F,mx[i][j][j1].S);
swap(mxnode[i][j][j1].F,mxnode[i][j][j1].S);
mx[i][j][j1].F = val;
mxnode[i][j][j1].F = x;
}
else
if(val>mx[i][j][j1].S){
mx[i][j][j1].S = val;
mxnode[i][j][j1].S = x;
}
}
}
for(auto x:adj[i]){
if(x==pre)continue;
sum -= a[x];
for(int j=1;j<=100;j++){
LL val[2] = {0,0};
for(int j1=0;j1<2;j1++)
if(mxnode[i][j][j1].F==x)
val[j1] = mx[i][j][j1].S;
else
val[j1] = mx[i][j][j1].F;
dp[i][j] = max(val[0],val[1]+sum);
}
rootshift(x,i);
sum += a[x];
}
}
int main(){
//freopen("out.txt","w",stdout);
scanf("%d%d",&n,&k);
for(int i=1;i<=n;i++)
scanf("%lld",&a[i]);
for(int i=1;i<n;i++){
int u,v;
scanf("%d%d",&u,&v);
adj[u].push_back(v);
adj[v].push_back(u);
}
dfs(1,0);
rootshift(1,0);
printf("%lld\n",ans);
}