#include <bits/stdc++.h>
#define int long long
// #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
// #pragma GCC optimize("O3")
// #pragma GCC optimize("unroll-loops")
#define F first
#define S second
#define mp make_pair
#define pb push_back
#define all(x) x.begin(), x.end()
#define kill(x) cout << x << "\n", exit(0)
#define pii pair<int, int>
#define endl "\n"
using namespace std;
typedef long long ll;
typedef long double ld;
const int MAXN = (int)1e5 + 7;
const int MOD = (int)1e9 + 7;
const int INF = (int)1e9 + 7;
const int SQ = 107;
int n, m, t, k, tmp, tmp2, tmp3, tmp4, flag, u, v, w, ans;
int arr[MAXN], dp[MAXN][SQ], dp2[MAXN][SQ], sum[MAXN];
int ps[SQ][2][MAXN], pd[MAXN][SQ];
vector<int> adj[MAXN], vec[MAXN];
void DFSsz(int v, int p=-1) {
sum[v] = 0; vec[v].clear();
for (int u:adj[v]) {
if (u == p) continue;
DFSsz(u, v);
sum[v] += arr[u]; vec[v].pb(u);
}
for (int j=1; j<=k; j++) {
for (int u:vec[v]) {
dp[v][j] = max({dp[v][j], max(dp[u][j], dp2[u][j])});
dp2[v][j] = max({dp2[v][j], max(dp[u][j-1], dp2[u][j-1])+sum[v]});
}
}
for (int j=1; j<=k; j++) {
ps[j][0][0] = ps[j][1][vec[v].size()+1] = 0;
for (int i=1; i<=vec[v].size(); i++) {
int u = vec[v][i-1];
ps[j][0][i] = max(ps[j][0][i-1], max(dp[u][j], dp2[u][j]));
}
for (int i=vec[v].size(); i>=1; i--) {
int u = vec[v][i-1];
ps[j][1][i] = max(ps[j][1][i+1], max(dp[u][j], dp2[u][j]));
}
for (int i=1; i<=vec[v].size(); i++) pd[vec[v][i-1]][j] = max(ps[j][0][i-1], ps[j][1][i+1]);
}
}
void Walk(int v, int p, vector<int> mul) {
for (int j=1; j<=k; j++) ans = max({ans, dp[v][j], dp2[v][j]+arr[p], mul[j], mul[j-1]+sum[v]+arr[p]});
for (int u:adj[v]) {
if (u == p) continue;
vector<int> mul2(k+1);
for (int j=1; j<=k; j++) mul2[j] = max({pd[u][j], mul[j], max(pd[u][j-1], mul[j-1])+sum[v]-arr[u]+arr[p]});
Walk(u, v, mul2);
}
}
int32_t main() {
#ifdef LOCAL
freopen("inp.in", "r", stdin);
freopen("res.out", "w", stdout);
#else
ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
#endif
cin >> n >> k;
for (int i=1; i<=n; i++) cin >> arr[i];
for (int i=1; i<n; i++) {
cin >> u >> v;
adj[u].pb(v); adj[v].pb(u);
}
DFSsz(1);
Walk(1, 0, vector<int>(k+1, -INF));
cout << ans << endl;
return 0;
}
Compilation message
chase.cpp: In function 'void DFSsz(long long int, long long int)':
chase.cpp:62:24: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
62 | for (int i=1; i<=vec[v].size(); i++) {
| ~^~~~~~~~~~~~~~~
chase.cpp:74:24: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
74 | for (int i=1; i<=vec[v].size(); i++) pd[vec[v][i-1]][j] = max(ps[j][0][i-1], ps[j][1][i+1]);
| ~^~~~~~~~~~~~~~~
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
2 ms |
16988 KB |
Output is correct |
2 |
Correct |
2 ms |
16988 KB |
Output is correct |
3 |
Correct |
2 ms |
14940 KB |
Output is correct |
4 |
Correct |
3 ms |
29356 KB |
Output is correct |
5 |
Correct |
2 ms |
10844 KB |
Output is correct |
6 |
Correct |
2 ms |
19036 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
2 ms |
16988 KB |
Output is correct |
2 |
Correct |
2 ms |
16988 KB |
Output is correct |
3 |
Correct |
2 ms |
14940 KB |
Output is correct |
4 |
Correct |
3 ms |
29356 KB |
Output is correct |
5 |
Correct |
2 ms |
10844 KB |
Output is correct |
6 |
Correct |
2 ms |
19036 KB |
Output is correct |
7 |
Correct |
11 ms |
99264 KB |
Output is correct |
8 |
Correct |
3 ms |
22104 KB |
Output is correct |
9 |
Correct |
3 ms |
21084 KB |
Output is correct |
10 |
Correct |
13 ms |
99008 KB |
Output is correct |
11 |
Correct |
7 ms |
61592 KB |
Output is correct |
12 |
Correct |
5 ms |
28784 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
347 ms |
357372 KB |
Output is correct |
2 |
Correct |
336 ms |
365012 KB |
Output is correct |
3 |
Correct |
278 ms |
258632 KB |
Output is correct |
4 |
Correct |
133 ms |
243024 KB |
Output is correct |
5 |
Correct |
359 ms |
332372 KB |
Output is correct |
6 |
Correct |
375 ms |
332116 KB |
Output is correct |
7 |
Correct |
432 ms |
327908 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
2 ms |
16988 KB |
Output is correct |
2 |
Correct |
2 ms |
16988 KB |
Output is correct |
3 |
Correct |
2 ms |
14940 KB |
Output is correct |
4 |
Correct |
3 ms |
29356 KB |
Output is correct |
5 |
Correct |
2 ms |
10844 KB |
Output is correct |
6 |
Correct |
2 ms |
19036 KB |
Output is correct |
7 |
Correct |
11 ms |
99264 KB |
Output is correct |
8 |
Correct |
3 ms |
22104 KB |
Output is correct |
9 |
Correct |
3 ms |
21084 KB |
Output is correct |
10 |
Correct |
13 ms |
99008 KB |
Output is correct |
11 |
Correct |
7 ms |
61592 KB |
Output is correct |
12 |
Correct |
5 ms |
28784 KB |
Output is correct |
13 |
Correct |
347 ms |
357372 KB |
Output is correct |
14 |
Correct |
336 ms |
365012 KB |
Output is correct |
15 |
Correct |
278 ms |
258632 KB |
Output is correct |
16 |
Correct |
133 ms |
243024 KB |
Output is correct |
17 |
Correct |
359 ms |
332372 KB |
Output is correct |
18 |
Correct |
375 ms |
332116 KB |
Output is correct |
19 |
Correct |
432 ms |
327908 KB |
Output is correct |
20 |
Correct |
395 ms |
330064 KB |
Output is correct |
21 |
Correct |
37 ms |
20560 KB |
Output is correct |
22 |
Correct |
370 ms |
327668 KB |
Output is correct |
23 |
Correct |
132 ms |
243028 KB |
Output is correct |
24 |
Correct |
381 ms |
327924 KB |
Output is correct |
25 |
Correct |
273 ms |
258116 KB |
Output is correct |
26 |
Correct |
333 ms |
360016 KB |
Output is correct |
27 |
Correct |
309 ms |
360016 KB |
Output is correct |