/** I can do this all day **/
#pragma GCC optimize("O2")
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef long double ld;
typedef pair<int,int> pii;
typedef pair<ll,ll> pll;
#define all(x) (x).begin(),(x).end()
#define F first
#define S second
#define Mp make_pair
#define SZ(x) (int)x.size()
#define fast_io ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);
#define file_io freopen("in.txt" , "r+" , stdin) ; freopen("out.txt" , "w+" , stdout);
const int N = 1e5 + 5;
const ll mod = 1e9 + 7;
const ll mod2 = 998244353;
const ll inf = 1e18;
const int LOG = 22;
const int M = 1e2 + 5;
int n, k, Par[N];
vector < int > G[N];
ll tot, A[N], nei[N], dp_down[2][M][N], dp_up[M][N];
void dfs_down(int v, int P)
{
Par[v] = P;
dp_up[1][v] = nei[v];
for(auto u : G[v])
{
if(u == P) continue;
dfs_down(u, v);
for(int x = 0; x <= k; x ++)
{
dp_up[x][v] = max(dp_up[x][v], dp_up[x][u]);
if(x)
{
dp_up[x][v] = max(dp_up[x][v], dp_up[x - 1][u] + nei[v] - A[u]);
}
}
}
}
void dfs(int v, int P)
{
for(int i = 1; i <= k; i ++)
{
dp_down[1][i][v] = nei[v] - A[P];
}
dp_down[1][0][v] = -inf;
for(auto u : G[v])
{
if(u == P) continue;
dfs(u, v);
for(int x = 0; x <= k; x ++)
{
tot = max(tot, dp_up[x][u] + dp_down[0][k - x][v]);
if(x != k) tot = max(tot, dp_up[x][u] + dp_down[1][k - x][v] + A[P] - A[u]);
///printf("x = %d v = %d u = %d cu0 = %lld cu1 = %lld\n", x, v, u, dp_up[x][u] + dp_down[0][k - x][v], dp_up[x][u] + dp_down[1][k - x][v] + A[P] - A[u]);
}
for(int x = 0; x <= k; x ++)
{
dp_down[0][x][v] = max(dp_down[0][x][v], max(dp_down[0][x][u], dp_down[1][x][u]));
if(x)
{
dp_down[1][x][v] = max(dp_down[1][x][v], nei[v] - A[P] + max(dp_down[1][x - 1][u], dp_down[0][x - 1][u]));
}
}
}
reverse(all(G[v]));
for(int i = 0; i < M; i ++) dp_down[0][i][v] = dp_down[1][i][v] = 0;
for(int i = 1; i <= k; i ++) dp_down[1][i][v] = nei[v] - A[P];
dp_down[1][0][v] = -inf;
for(auto u : G[v])
{
if(u == P) continue;
///dfs(u, v);
for(int x = 0; x <= k; x ++)
{
tot = max(tot, dp_up[x][u] + dp_down[0][k - x][v]);
if(x != k) tot = max(tot, dp_up[x][u] + dp_down[1][k - x][v] + A[P] - A[u]);
}
for(int x = 0; x <= k; x ++)
{
dp_down[0][x][v] = max(dp_down[0][x][v], max(dp_down[0][x][u], dp_down[1][x][u]));
if(x)
{
dp_down[1][x][v] = max(dp_down[1][x][v], nei[v] - A[P] + max(dp_down[1][x - 1][u], dp_down[0][x - 1][u]));
}
}
}
}
int main()
{
scanf("%d%d", &n, &k);
for(int i = 1; i <= n; i ++)
{
scanf("%lld", &A[i]);
}
for(int i = 1; i < n; i ++)
{
int a, b;
scanf("%d%d", &a, &b);
G[a].push_back(b);
G[b].push_back(a);
nei[a] += A[b];
nei[b] += A[a];
}
dfs_down(1, 0);
dfs(1, 0);
for(int i = 1; i <= n; i ++)
{
tot = max(tot, dp_down[1][k][i] + A[Par[i]]);
tot = max(tot, dp_up[k][i]);
}
/*for(int i = 1; i <= n; i ++)
{
for(int j = 1; j <= k; j ++)
{
printf("i = %d j = %d dp = %lld\n", i, j, dp_up[j][i]);
}
}
*/
printf("%lld", tot);
return 0;
}
/** test corner cases(n = 1?) watch for overflow or minus indices **/
Compilation message
chase.cpp: In function 'int main()':
chase.cpp:105:7: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
105 | scanf("%d%d", &n, &k);
| ~~~~~^~~~~~~~~~~~~~~~
chase.cpp:108:8: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
108 | scanf("%lld", &A[i]);
| ~~~~~^~~~~~~~~~~~~~~
chase.cpp:113:8: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
113 | scanf("%d%d", &a, &b);
| ~~~~~^~~~~~~~~~~~~~~~
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
3 ms |
3788 KB |
Output is correct |
2 |
Correct |
2 ms |
3788 KB |
Output is correct |
3 |
Correct |
2 ms |
3788 KB |
Output is correct |
4 |
Correct |
2 ms |
3788 KB |
Output is correct |
5 |
Correct |
3 ms |
3788 KB |
Output is correct |
6 |
Correct |
2 ms |
3788 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
3 ms |
3788 KB |
Output is correct |
2 |
Correct |
2 ms |
3788 KB |
Output is correct |
3 |
Correct |
2 ms |
3788 KB |
Output is correct |
4 |
Correct |
2 ms |
3788 KB |
Output is correct |
5 |
Correct |
3 ms |
3788 KB |
Output is correct |
6 |
Correct |
2 ms |
3788 KB |
Output is correct |
7 |
Correct |
7 ms |
6476 KB |
Output is correct |
8 |
Correct |
4 ms |
5524 KB |
Output is correct |
9 |
Correct |
4 ms |
5452 KB |
Output is correct |
10 |
Correct |
7 ms |
6732 KB |
Output is correct |
11 |
Correct |
5 ms |
5836 KB |
Output is correct |
12 |
Correct |
4 ms |
5580 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
856 ms |
217992 KB |
Output is correct |
2 |
Correct |
873 ms |
217928 KB |
Output is correct |
3 |
Correct |
463 ms |
174252 KB |
Output is correct |
4 |
Correct |
541 ms |
173856 KB |
Output is correct |
5 |
Correct |
1499 ms |
250508 KB |
Output is correct |
6 |
Correct |
1504 ms |
251440 KB |
Output is correct |
7 |
Correct |
1497 ms |
251172 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
3 ms |
3788 KB |
Output is correct |
2 |
Correct |
2 ms |
3788 KB |
Output is correct |
3 |
Correct |
2 ms |
3788 KB |
Output is correct |
4 |
Correct |
2 ms |
3788 KB |
Output is correct |
5 |
Correct |
3 ms |
3788 KB |
Output is correct |
6 |
Correct |
2 ms |
3788 KB |
Output is correct |
7 |
Correct |
7 ms |
6476 KB |
Output is correct |
8 |
Correct |
4 ms |
5524 KB |
Output is correct |
9 |
Correct |
4 ms |
5452 KB |
Output is correct |
10 |
Correct |
7 ms |
6732 KB |
Output is correct |
11 |
Correct |
5 ms |
5836 KB |
Output is correct |
12 |
Correct |
4 ms |
5580 KB |
Output is correct |
13 |
Correct |
856 ms |
217992 KB |
Output is correct |
14 |
Correct |
873 ms |
217928 KB |
Output is correct |
15 |
Correct |
463 ms |
174252 KB |
Output is correct |
16 |
Correct |
541 ms |
173856 KB |
Output is correct |
17 |
Correct |
1499 ms |
250508 KB |
Output is correct |
18 |
Correct |
1504 ms |
251440 KB |
Output is correct |
19 |
Correct |
1497 ms |
251172 KB |
Output is correct |
20 |
Correct |
1490 ms |
251416 KB |
Output is correct |
21 |
Correct |
512 ms |
173892 KB |
Output is correct |
22 |
Correct |
1489 ms |
251384 KB |
Output is correct |
23 |
Correct |
531 ms |
173764 KB |
Output is correct |
24 |
Correct |
1481 ms |
251360 KB |
Output is correct |
25 |
Correct |
433 ms |
173756 KB |
Output is correct |
26 |
Correct |
859 ms |
217960 KB |
Output is correct |
27 |
Correct |
858 ms |
218140 KB |
Output is correct |