#include <bits/stdc++.h>
using namespace std;
#define all(x) x.begin(), x.end()
#define ar array
#define pb push_back
#define ln '\n'
#define int long long
using i64 = long long;
template <class F, class _S>
bool chmin(F &u, const _S &v){
bool flag = false;
if ( u > v ){
u = v; flag |= true;
}
return flag;
}
template <class F, class _S>
bool chmax(F &u, const _S &v){
bool flag = false;
if ( u < v ){
u = v; flag |= true;
}
return flag;
}
signed main(){
ios_base::sync_with_stdio(false);
cin.tie(nullptr);
int n, k; cin >> n >> k;
vector <int> a(n);
for ( auto &u: a ) cin >> u;
vector <vector<int>> G(n);
for ( int i = 0; i + 1 < n; i++ ){
int u, v; cin >> u >> v;
--u, --v;
G[u].pb(v), G[v].pb(u);
}
vector <int> s(n);
for ( int i = 0; i < n; i++ ){
for ( auto &v: G[i] ){
s[i] += a[v];
}
}
vector <vector<multiset<int>>> dp(n, vector <multiset<int>> (k + 1));
for ( int i = 0; i < n; i++ ){
for ( int j = 0; j <= k; j++ ){
dp[i][j].insert(0);
}
}
auto er = [&](auto &s, int u){
s.erase(s.find(u));
};
auto get = [&](int v, int j, int u){
return max(*dp[v][j].rbegin(), (j > 0 ? *dp[v][j - 1].rbegin() + s[v] - a[u] : 0));
};
auto dfs = [&](auto dfs, int u, int p) -> void{
dp[u][0].insert(0);
for ( auto &v: G[u] ){
if ( v != p ){
dfs(dfs, v, u);
for ( int j = 0; j <= k; j++ ){
dp[u][j].insert(get(v, j, u));
}
}
}
};
dfs(dfs, 0, -1);
int opt = 0;
auto dfs2 = [&](auto dfs2, int u, int p) -> void{
for ( int j = 1; j <= k; j++ ){
chmax(opt, *dp[u][j - 1].rbegin() + s[u]);
}
for ( auto &v: G[u] ){
if ( v == p ) continue;
for ( int j = 0; j <= k; j++ ){
er(dp[u][j], get(v, j, u));
}
for ( int j = 0; j <= k; j++ ){
dp[v][j].insert(get(u, j, v));
}
dfs2(dfs2, v, u);
for ( int j = 0; j <= k; j++ ){
er(dp[v][j], get(u, j, v));
}
for ( int j = 0; j <= k; j++ ){
dp[u][j].insert(get(v, j, u));
}
}
};
dfs2(dfs2, 0, -1);
cout << opt;
cout << '\n';
}
/*
12 2
2 3 3 8 1 5 6 7 8 3 5 4
2 1
2 7
3 4
4 7
7 6
5 6
6 8
6 9
7 10
10 11
10 12
5 2
3 1 2 4 1
5 2
1 2
2 3
3 4
*/
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
0 ms |
344 KB |
Output is correct |
2 |
Correct |
0 ms |
344 KB |
Output is correct |
3 |
Correct |
0 ms |
348 KB |
Output is correct |
4 |
Correct |
0 ms |
600 KB |
Output is correct |
5 |
Correct |
0 ms |
348 KB |
Output is correct |
6 |
Correct |
0 ms |
348 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
0 ms |
344 KB |
Output is correct |
2 |
Correct |
0 ms |
344 KB |
Output is correct |
3 |
Correct |
0 ms |
348 KB |
Output is correct |
4 |
Correct |
0 ms |
600 KB |
Output is correct |
5 |
Correct |
0 ms |
348 KB |
Output is correct |
6 |
Correct |
0 ms |
348 KB |
Output is correct |
7 |
Correct |
25 ms |
14940 KB |
Output is correct |
8 |
Correct |
2 ms |
1372 KB |
Output is correct |
9 |
Correct |
2 ms |
1372 KB |
Output is correct |
10 |
Correct |
24 ms |
14848 KB |
Output is correct |
11 |
Correct |
12 ms |
4980 KB |
Output is correct |
12 |
Correct |
3 ms |
1884 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Runtime error |
330 ms |
524288 KB |
Execution killed with signal 9 |
2 |
Halted |
0 ms |
0 KB |
- |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
0 ms |
344 KB |
Output is correct |
2 |
Correct |
0 ms |
344 KB |
Output is correct |
3 |
Correct |
0 ms |
348 KB |
Output is correct |
4 |
Correct |
0 ms |
600 KB |
Output is correct |
5 |
Correct |
0 ms |
348 KB |
Output is correct |
6 |
Correct |
0 ms |
348 KB |
Output is correct |
7 |
Correct |
25 ms |
14940 KB |
Output is correct |
8 |
Correct |
2 ms |
1372 KB |
Output is correct |
9 |
Correct |
2 ms |
1372 KB |
Output is correct |
10 |
Correct |
24 ms |
14848 KB |
Output is correct |
11 |
Correct |
12 ms |
4980 KB |
Output is correct |
12 |
Correct |
3 ms |
1884 KB |
Output is correct |
13 |
Runtime error |
330 ms |
524288 KB |
Execution killed with signal 9 |
14 |
Halted |
0 ms |
0 KB |
- |