#include <bits/stdc++.h>
#define DIM 100002
#define INF 2000000000000000000LL
using namespace std;
vector <int> L[DIM];
int v[DIM],fth[DIM];
long long sum[DIM],maxi;
int n,x,y,i,j,nr,k;
long long dp_up[DIM][102][2], dp_down[DIM][102][2];
/// dp_up[nod][i] = care e suma maxima daca incep undeva in subarborele lui nod,
/// ajung in nod si am consumat maxim i boabe
/// dp_down[nod][i] = la fel doar ca incep din nod si ma duc in jos
void dfs (int nod, int tata){
int ok = 0;
fth[nod] = tata;
for (auto vecin : L[nod])
if (vecin != tata){
ok = 1;
dfs (vecin,nod);
}
dp_down[nod][1][1] = dp_up[nod][1][1] = sum[nod];
for (auto vecin : L[nod]){
if (vecin == tata)
continue;
/// dp_down + dp_up
for (int i=0;i<=nr;i++){
maxi = max (maxi,dp_down[nod][i][0] + dp_up[vecin][nr-i][0]);
maxi = max (maxi,dp_down[nod][i][0] + dp_up[vecin][nr-i][1]);
/// daca as pune in nod
maxi = max (maxi,dp_down[nod][i][1] + dp_up[vecin][nr-i][0] - v[vecin]);
maxi = max (maxi,dp_down[nod][i][1] + dp_up[vecin][nr-i][1] - v[vecin]);
/// acum fac pe invers
maxi = max (maxi,dp_down[vecin][i][0] + dp_up[nod][nr-i][0]);
maxi = max (maxi,dp_down[vecin][i][0] + dp_up[nod][nr-i][1]);
maxi = max (maxi,dp_down[vecin][i][1] + dp_up[nod][nr-i][0] - v[nod]);
maxi = max (maxi,dp_down[vecin][i][1] + dp_up[nod][nr-i][1] - v[nod]);
}
for (int i=1;i<=nr;i++){
/// vin de jos, nu mai conteaza nimic din ce am in nod, pt ca nu influenteaza
dp_up[nod][i][0] = max (dp_up[nod][i][0], dp_up[nod][i-1][0]);
dp_up[nod][i][0] = max (dp_up[nod][i][0], max(dp_up[vecin][i][0],dp_up[vecin][i][1]));
/// cand pun boaba in i trb sa am grija sa scad valoarea din fiul in care ma aflu
dp_up[nod][i][1] = max (dp_up[nod][i][1], dp_up[nod][i-1][1]);
dp_up[nod][i][1] = max (dp_up[nod][i][1], dp_up[vecin][i-1][0] + sum[nod] - v[vecin]);
dp_up[nod][i][1] = max (dp_up[nod][i][1], dp_up[vecin][i-1][1] + sum[nod] - v[vecin]);
/// acum ma duc spre vecin
dp_down[nod][i][0] = max (dp_down[nod][i][0], dp_down[nod][i-1][0]);
dp_down[nod][i][0] = max (dp_down[nod][i][0], dp_down[vecin][i][0]);
dp_down[nod][i][0] = max (dp_down[nod][i][0], dp_down[vecin][i][1] - v[nod]);
dp_down[nod][i][1] = max (dp_down[nod][i][1], dp_down[nod][i-1][1]);
dp_down[nod][i][1] = max (dp_down[nod][i][1], dp_down[vecin][i-1][0] + sum[nod]);
dp_down[nod][i][1] = max (dp_down[nod][i][1], dp_down[vecin][i-1][1] + sum[nod] - v[vecin]);
}
}
}
int main (){
//ifstream cin ("date.in");
//ofstream cout ("date.out");
cin>>n>>nr;
for (i=1;i<=n;i++)
cin>>v[i];
for (i=1;i<n;i++){
cin>>x>>y;
L[x].push_back(y);
L[y].push_back(x);
}
for (i=1;i<=n;i++){
//sum[i] = v[i];
for (auto it : L[i])
sum[i] += v[it];
}
if (nr == 0){
cout<<0;
return 0;
}
if (nr == 1){
long long maxi = 0;
for (i=1;i<=n;i++)
maxi = max (maxi,sum[i] - v[i]);
cout<<maxi;
return 0;
}
for (i=1;i<=n;i++)
for (j=1;j<=nr;j++){
dp_down[i][j][0] = dp_down[i][j][1] = -INF;
dp_up[i][j][0] = dp_up[i][j][1] = -INF;
}
dfs (1,0);
for (i=1;i<=n;i++)
for (j=1;j<=nr;j++){
maxi = max (maxi,max(dp_down[i][j][0],dp_down[i][j][1]));
maxi = max (maxi,max(dp_up[i][j][0],dp_up[i][j][1]));
}
cout<<maxi;
return 0;
}
Compilation message
chase.cpp: In function 'void dfs(int, int)':
chase.cpp:16:9: warning: variable 'ok' set but not used [-Wunused-but-set-variable]
int ok = 0;
^~
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
6 ms |
2816 KB |
Output is correct |
| 2 |
Correct |
6 ms |
2688 KB |
Output is correct |
| 3 |
Incorrect |
6 ms |
2688 KB |
Output isn't correct |
| 4 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
6 ms |
2816 KB |
Output is correct |
| 2 |
Correct |
6 ms |
2688 KB |
Output is correct |
| 3 |
Incorrect |
6 ms |
2688 KB |
Output isn't correct |
| 4 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
538 ms |
336248 KB |
Output is correct |
| 2 |
Correct |
530 ms |
336888 KB |
Output is correct |
| 3 |
Incorrect |
500 ms |
329456 KB |
Output isn't correct |
| 4 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
6 ms |
2816 KB |
Output is correct |
| 2 |
Correct |
6 ms |
2688 KB |
Output is correct |
| 3 |
Incorrect |
6 ms |
2688 KB |
Output isn't correct |
| 4 |
Halted |
0 ms |
0 KB |
- |
