#include <bits/stdc++.h>
using namespace std;
using ll = long long;
#define int ll
const int N = 1e5 + 5, V = 105, inf = 1e9;
int dp1[V][N], dp2[V][N], p[N], v, ans; //<- sper ca am destula memorie
vector<int> vec[N];
void dfs(int nod, int papa) {
for (auto i : vec[nod])
if (i != papa)
dfs(i, nod);
int sumvec = p[papa], mx = -inf;
vector<vector<int>> dpp(2, vector<int>(v + 1, -inf));
for (auto i : vec[nod])
if (i != papa) {
mx = max(mx, dp2[v - 1][i] - p[i]);
sumvec += p[i];
for (int j = 0; j <= v; j ++) {
if (dpp[0][j] < dp1[j][i] - p[i]) {
dpp[1][j] = dpp[0][j];
dpp[0][j] = dp1[j][i] - p[i];
} else if (dpp[1][j] < dp1[j][i] - p[i])
dpp[1][j] = dp1[j][i] - p[i];
}
}
/*
avem 3 variante:
incepe din nod, se termina in nod sau contine nod undeva random
*/
ans = max(ans, dpp[0][v - 1] + sumvec);
ans = max(ans, mx + sumvec + p[nod]);
for (auto i : vec[nod]) {
if (i == papa)
continue;
for (int j = 1; j < v - 1; j ++) {
int x = v - 1 - j;
ans = max(ans, dp2[j][i] - p[i] + p[nod] + sumvec + ((dp1[x][i] - p[i] == dpp[0][x]) ? dpp[1][x] : dpp[0][x]));
}
}
for (auto i : vec[nod]) {
if (i != papa) {
for (int j = 1; j < v; j ++)
dp1[j + 1][nod] = max(dp1[j + 1][nod], dp1[j][i] + sumvec - p[papa] + p[nod]);
for (int j = 1; j < v; j ++)
dp2[j + 1][nod] = max(dp2[j + 1][nod], dp2[j][i] + sumvec - p[papa] + p[nod]);
}
}
dp1[1][nod] = sumvec - p[papa] + p[nod];
dp2[1][nod] = sumvec - p[papa];
for (int i = 1; i <= v; i ++) {
dp1[i][nod] = max(dp1[i][nod], dp1[i - 1][nod]);
dp2[i][nod] = max(dp2[i][nod], dp2[i - 1][nod]);
}
}
signed main() {
int n;
cin >> n >> v;
for (int i = 1; i <= n; i ++)
cin >> p[i];
memset(dp1, -inf, sizeof(dp1));
memset(dp2, -inf, sizeof(dp2));
for (int i = 1; i < n; i ++) {
int x, y;
cin >> x >> y;
vec[x].push_back(y);
vec[y].push_back(x);
}
dfs(1, 0);
cout << ans << '\n';
}
# | 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... |