This submission is migrated from previous version of oj.uz, which used different machine for grading. This submission may have different result if resubmitted.
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef long double ld;
const int N = 100100;
ll dp[N][110], dp2[N][110], a[N];
vector <int> g[N];
int n, k;
ll ans = 0;
void dfs(int v, int p){
ll mx1[101] = {0}, mx2[101] = {0};
ll sum = 0;
for (int i = 0; i < g[v].size(); i++){
sum += a[g[v][i]];
}
dp[v][0] = dp2[v][0] = 0;
dp2[v][1] = sum;
for (int i = 0; i < g[v].size(); i++){
int to = g[v][i];
if (to == p)
continue;
dfs(to, v);
for (int j = 1; j <= k; j++){
dp[v][j] = max(dp[v][j], max(dp[to][j], dp[to][j-1] + sum - a[p]));
dp2[v][j] = max(dp2[v][j], max(dp2[to][j], dp2[to][j-1] + sum - a[to]));
if (dp[to][j] >= mx1[j]){
mx2[j] = mx1[j];
mx1[j] = dp[to][j];
} else {
mx2[j] = max(mx2[j], dp[to][j]);
}
}
}
for (int i = 0; i < g[v].size(); i++){
int to = g[v][i];
if (to == p)
continue;
ans = max(ans, sum + dp[to][k-1]);
for (int j = 0; j <= k; j++){
ll mx = max(dp2[to][j], (j == 0 ? 0 : dp2[to][j-1] + sum - a[to]));
if (mx1[k-j] == dp[to][k-j]){
ans = max(ans, mx + mx2[k-j]);
} else {
ans = max(ans, mx + mx1[k-j]);
}
}
}
}
int main(){
ios_base::sync_with_stdio(0);cin.tie(0);
cin >> n >> k;
for (int i = 1; i <= n; i++){
cin >> a[i];
}
for (int i = 2; i <= n; i++){
int u, v;
cin >> u >> v;
g[u].push_back(v);
g[v].push_back(u);
}
if (k == 0)
return cout << 0, 0;
dfs(1, 0);
cout << ans;
}
/**
*/
Compilation message (stderr)
chase.cpp: In function 'void dfs(int, int)':
chase.cpp:21:23: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
for (int i = 0; i < g[v].size(); i++){
~~^~~~~~~~~~~~~
chase.cpp:26:23: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
for (int i = 0; i < g[v].size(); i++){
~~^~~~~~~~~~~~~
chase.cpp:42:23: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
for (int i = 0; i < g[v].size(); i++){
~~^~~~~~~~~~~~~
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |