#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];
void init2() {
for (int i = 0; i <= n; i++)
for (int c = 0; c <= 1; c++)
dp2[i][c].resize(adj[i].size() + 2);
}
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 = 1; i <= n; i++) {
int u = rev_ind[i];
multiset<lint> trans;
trans.insert(0);
for (int v : children[u]) trans.insert(dp1[v][opp_c]); // Should have 0
for (int j = 0; j < children[u].size(); j++) {
if (c == 0) continue;
int v = children[u][j];
lint leave = dp2[par[u]][p][child_ind[u]];
lint take = dp2[par[u]][opp_p][child_ind[u]] + val_sum[u] - val[v];
dp2[u][p][j] = max(take, leave);
}
for (int j = 0; j < children[u].size(); j++) {
int v = children[u][j];
trans.erase(trans.find(dp1[v][opp_c]));
assert(trans.size());
ans = max(ans, dp2[u][p][j] + *trans.rbegin());
trans.insert(dp1[v][opp_c]);
}
}
}
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:101:31: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
101 | for (int j = 0; j < children[u].size(); j++) {
| ~~^~~~~~~~~~~~~~~~~~~~
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
3 ms |
16988 KB |
Output is correct |
2 |
Incorrect |
3 ms |
16988 KB |
Output isn't correct |
3 |
Halted |
0 ms |
0 KB |
- |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
3 ms |
16988 KB |
Output is correct |
2 |
Incorrect |
3 ms |
16988 KB |
Output isn't correct |
3 |
Halted |
0 ms |
0 KB |
- |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Incorrect |
3911 ms |
115244 KB |
Output isn't correct |
2 |
Halted |
0 ms |
0 KB |
- |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
3 ms |
16988 KB |
Output is correct |
2 |
Incorrect |
3 ms |
16988 KB |
Output isn't correct |
3 |
Halted |
0 ms |
0 KB |
- |