#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]
void init() {
for (int i = 1; i <= n; i++) {
adj[i].push_back(0);
adj_ind[i].push_back(0);
}
}
vector<int> seen[MAX_N];
vector<pii> topo_order;
void dfs(int u, int prev) {
seen[u][prev] = 1;
for (int i = 1; i < adj[u].size(); i++) {
if (i == prev) continue;
int v = adj[u][i];
if (seen[v][adj_ind[u][i]] == 1) assert(true == false);
if (!seen[v][adj_ind[u][i]]) dfs(v, adj_ind[u][i]);
}
seen[u][prev] = 2;
topo_order.push_back({u, prev});
}
vector<lint> dp[MAX_N][MAX_K];
void precomp() {
for (int i = 1; i <= n; i++)
seen[i].resize(adj[i].size() + 2);
for (int i = 1; i <= n; i++)
dfs(i, 0);
reverse(topo_order.begin(), topo_order.end());
assert(topo_order.size() <= 3 * 1e5);
for (int i = 1; i <= n; i++) {
for (int c = 0; c <= k; c++) {
dp[i][c].resize(adj[i].size() + 2);
}
}
}
int main() {
// freopen("chase.in", "r", stdin);
cin >> n >> k;
init();
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);
int u_ind = adj[u].size() - 1, v_ind = adj[v].size() - 1;
adj_ind[u].push_back(v_ind);
adj_ind[v].push_back(u_ind);
}
precomp();
for (int c = 1; c <= k; c++) {
for (int j = topo_order.size() - 1; j >= 0; j--) {
int u = topo_order[j].first, prev = topo_order[j].second;
lint leave = 0;
for (int i = 1; i < adj[u].size(); i++)
if (i != prev) leave = max(leave, dp[adj[u][i]][c][adj_ind[u][i]]);
lint take = 0;
for (int i = 1; i < adj[u].size(); i++)
if (i != prev) take = max(take, dp[adj[u][i]][c - 1][adj_ind[u][i]]);
for (int i = 1; i < adj[u].size(); i++)
if (i != prev) take += val[adj[u][i]];
dp[u][c][prev] = max(take, leave);
}
}
lint ans = 0;
for (int i = 1; i <= n; i++) ans = max(ans, dp[i][k][0]);
cout << ans << '\n';
}
Compilation message
chase.cpp: In function 'void dfs(int, int)':
chase.cpp:21:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
21 | for (int i = 1; i < adj[u].size(); i++) {
| ~~^~~~~~~~~~~~~~~
chase.cpp: In function 'int main()':
chase.cpp:72:31: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
72 | for (int i = 1; i < adj[u].size(); i++)
| ~~^~~~~~~~~~~~~~~
chase.cpp:76:31: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
76 | for (int i = 1; i < adj[u].size(); i++)
| ~~^~~~~~~~~~~~~~~
chase.cpp:78:31: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
78 | for (int i = 1; i < adj[u].size(); i++)
| ~~^~~~~~~~~~~~~~~
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
58 ms |
254036 KB |
Output is correct |
2 |
Correct |
56 ms |
254036 KB |
Output is correct |
3 |
Correct |
59 ms |
254036 KB |
Output is correct |
4 |
Correct |
61 ms |
254012 KB |
Output is correct |
5 |
Correct |
62 ms |
254020 KB |
Output is correct |
6 |
Correct |
60 ms |
254056 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
58 ms |
254036 KB |
Output is correct |
2 |
Correct |
56 ms |
254036 KB |
Output is correct |
3 |
Correct |
59 ms |
254036 KB |
Output is correct |
4 |
Correct |
61 ms |
254012 KB |
Output is correct |
5 |
Correct |
62 ms |
254020 KB |
Output is correct |
6 |
Correct |
60 ms |
254056 KB |
Output is correct |
7 |
Correct |
73 ms |
259412 KB |
Output is correct |
8 |
Correct |
60 ms |
254548 KB |
Output is correct |
9 |
Correct |
82 ms |
254520 KB |
Output is correct |
10 |
Correct |
80 ms |
259472 KB |
Output is correct |
11 |
Correct |
63 ms |
255568 KB |
Output is correct |
12 |
Correct |
61 ms |
254484 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Runtime error |
348 ms |
524288 KB |
Execution killed with signal 9 |
2 |
Halted |
0 ms |
0 KB |
- |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
58 ms |
254036 KB |
Output is correct |
2 |
Correct |
56 ms |
254036 KB |
Output is correct |
3 |
Correct |
59 ms |
254036 KB |
Output is correct |
4 |
Correct |
61 ms |
254012 KB |
Output is correct |
5 |
Correct |
62 ms |
254020 KB |
Output is correct |
6 |
Correct |
60 ms |
254056 KB |
Output is correct |
7 |
Correct |
73 ms |
259412 KB |
Output is correct |
8 |
Correct |
60 ms |
254548 KB |
Output is correct |
9 |
Correct |
82 ms |
254520 KB |
Output is correct |
10 |
Correct |
80 ms |
259472 KB |
Output is correct |
11 |
Correct |
63 ms |
255568 KB |
Output is correct |
12 |
Correct |
61 ms |
254484 KB |
Output is correct |
13 |
Runtime error |
348 ms |
524288 KB |
Execution killed with signal 9 |
14 |
Halted |
0 ms |
0 KB |
- |