#include <bits/stdc++.h>
using namespace std;
using pii = pair<int, int>;
using lint = long long;
const int MAX_N = 1e5 + 5, MAX_K = 1e2 + 5;
int n, k;
lint val[MAX_N];
vector<int> adj[MAX_N], adj_ind[MAX_N]; // adj_ind[u][i] = ind of adj list u is for adj[u][i]
int par[MAX_N];
vector<int> children[MAX_N];
void init1() {
for (int i = 1; i <= n; i++) children[i].push_back(0);
}
int child_ind[MAX_N];
int n_inds, ind[MAX_N], rev_ind[MAX_N];
void dfs(int u) {
ind[u] = ++n_inds;
rev_ind[ind[u]] = u;
for (int v : adj[u]) {
if (ind[v]) continue;
par[v] = u;
children[u].push_back(v);
child_ind[v] = children[u].size() - 1;
dfs(v);
}
}
lint val_sum[MAX_N];
void precomp() {
init1();
dfs(1);
for (int i = 1; i <= n; i++)
for (int j : adj[i]) val_sum[i] += val[j];
}
lint dp1[MAX_N][MAX_K];
vector<lint> dp2[MAX_N][2];
lint max_dp2[MAX_N][2];
void init2() {
for (int i = 1; i <= n; i++)
for (int c = 0; c <= 1; c++)
dp2[i][c].resize(adj[i].size() + 2);
}
lint pref_max[MAX_N], suff_max[MAX_N];
int main() {
// freopen("chase.in", "r", stdin);
cin >> n >> k;
for (int i = 1; i <= n; i++) cin >> val[i];
for (int i = 1; i < n; i++) {
int u, v; cin >> u >> v;
adj[u].push_back(v);
adj[v].push_back(u);
}
precomp();
for (int c = 1; c <= k; c++) {
for (int i = n; i >= 1; i--) {
int u = rev_ind[i];
lint leave = 0;
for (int v : children[u])
if (v != 0) leave = max(leave, dp1[v][c]);
lint take = 0;
for (int v : children[u])
if (v != 0) take = max(take, dp1[v][c - 1]);
take += val_sum[u] - val[par[u]];
dp1[u][c] = max(take, leave);
}
}
init2();
lint ans = 0;
for (int c = 0; c <= k; c++) {
int opp_c = k - c;
int p = c % 2, opp_p = (c + 1) % 2;
for (int i = n; i >= 1; i--) {
int u = rev_ind[i];
max_dp2[u][p] = 0;
for (int j = 0; j < children[u].size(); j++) {
pref_max[j] = max(((j == 0) ? 0 : pref_max[j - 1]), dp1[children[u][j]][opp_c]);
if (c == 0) continue;
int v = children[u][j];
lint leave = (v == 0) ? 0 : max_dp2[v][p];
lint take = val_sum[u] - val[v] + ((v == 0) ? 0 : max_dp2[v][opp_p]);
dp2[u][p][j] = max(take, leave);
max_dp2[u][p] = max(max_dp2[u][p], dp2[u][p][j]);
}
for (int j = children[u].size() - 1; j >= 0; j--)
suff_max[j] = max((j == children[u].size() - 1) ? 0 : suff_max[j + 1], dp1[children[u][j]][opp_c]);
for (int j = 0; j < children[u].size(); j++) {
lint tot_max = (j == 0) ? 0 : pref_max[j - 1];
tot_max = max(tot_max, ((j == children[u].size() - 1) ? 0 : suff_max[j + 1]));
ans = max(ans, dp2[u][p][j] + tot_max);
}
}
}
cout << ans << '\n';
}
Compilation message
chase.cpp: In function 'int main()':
chase.cpp:92:31: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
92 | for (int j = 0; j < children[u].size(); j++) {
| ~~^~~~~~~~~~~~~~~~~~~~
chase.cpp:104:38: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
104 | suff_max[j] = max((j == children[u].size() - 1) ? 0 : suff_max[j + 1], dp1[children[u][j]][opp_c]);
| ~~^~~~~~~~~~~~~~~~~~~~~~~~~
chase.cpp:106:31: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
106 | for (int j = 0; j < children[u].size(); j++) {
| ~~^~~~~~~~~~~~~~~~~~~~
chase.cpp:108:44: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
108 | tot_max = max(tot_max, ((j == children[u].size() - 1) ? 0 : suff_max[j + 1]));
| ~~^~~~~~~~~~~~~~~~~~~~~~~~~
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
4 ms |
17500 KB |
Output is correct |
2 |
Correct |
4 ms |
17580 KB |
Output is correct |
3 |
Correct |
4 ms |
17500 KB |
Output is correct |
4 |
Correct |
4 ms |
17500 KB |
Output is correct |
5 |
Correct |
3 ms |
17500 KB |
Output is correct |
6 |
Correct |
4 ms |
17496 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
4 ms |
17500 KB |
Output is correct |
2 |
Correct |
4 ms |
17580 KB |
Output is correct |
3 |
Correct |
4 ms |
17500 KB |
Output is correct |
4 |
Correct |
4 ms |
17500 KB |
Output is correct |
5 |
Correct |
3 ms |
17500 KB |
Output is correct |
6 |
Correct |
4 ms |
17496 KB |
Output is correct |
7 |
Correct |
8 ms |
19840 KB |
Output is correct |
8 |
Correct |
5 ms |
19956 KB |
Output is correct |
9 |
Correct |
5 ms |
19804 KB |
Output is correct |
10 |
Correct |
12 ms |
19804 KB |
Output is correct |
11 |
Correct |
6 ms |
19804 KB |
Output is correct |
12 |
Correct |
5 ms |
19804 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
2056 ms |
118204 KB |
Output is correct |
2 |
Correct |
1951 ms |
118184 KB |
Output is correct |
3 |
Correct |
669 ms |
115504 KB |
Output is correct |
4 |
Correct |
159 ms |
115272 KB |
Output is correct |
5 |
Correct |
2812 ms |
115272 KB |
Output is correct |
6 |
Correct |
2850 ms |
115268 KB |
Output is correct |
7 |
Correct |
3047 ms |
115264 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
4 ms |
17500 KB |
Output is correct |
2 |
Correct |
4 ms |
17580 KB |
Output is correct |
3 |
Correct |
4 ms |
17500 KB |
Output is correct |
4 |
Correct |
4 ms |
17500 KB |
Output is correct |
5 |
Correct |
3 ms |
17500 KB |
Output is correct |
6 |
Correct |
4 ms |
17496 KB |
Output is correct |
7 |
Correct |
8 ms |
19840 KB |
Output is correct |
8 |
Correct |
5 ms |
19956 KB |
Output is correct |
9 |
Correct |
5 ms |
19804 KB |
Output is correct |
10 |
Correct |
12 ms |
19804 KB |
Output is correct |
11 |
Correct |
6 ms |
19804 KB |
Output is correct |
12 |
Correct |
5 ms |
19804 KB |
Output is correct |
13 |
Correct |
2056 ms |
118204 KB |
Output is correct |
14 |
Correct |
1951 ms |
118184 KB |
Output is correct |
15 |
Correct |
669 ms |
115504 KB |
Output is correct |
16 |
Correct |
159 ms |
115272 KB |
Output is correct |
17 |
Correct |
2812 ms |
115272 KB |
Output is correct |
18 |
Correct |
2850 ms |
115268 KB |
Output is correct |
19 |
Correct |
3047 ms |
115264 KB |
Output is correct |
20 |
Correct |
3234 ms |
117364 KB |
Output is correct |
21 |
Correct |
120 ms |
35668 KB |
Output is correct |
22 |
Correct |
2608 ms |
117360 KB |
Output is correct |
23 |
Correct |
180 ms |
117244 KB |
Output is correct |
24 |
Correct |
2943 ms |
117360 KB |
Output is correct |
25 |
Correct |
649 ms |
117240 KB |
Output is correct |
26 |
Correct |
1919 ms |
120336 KB |
Output is correct |
27 |
Correct |
1882 ms |
120336 KB |
Output is correct |