#include <algorithm>
#include <iostream>
#include <vector>
using namespace std;
const int N = 100000;
const int K = 100;
const long long INF = 0x3f3f3f3f3f3f3f3fLL;
int aa[N], *ej[N], eo[N], k_;
long long bb[N], dp[N][K + 1][2], dq[N][K + 1], xx[K + 1];
void append(int i, int j) {
int o = eo[i]++;
if (!o)
ej[i] = (int *) malloc(sizeof *ej[i]);
else if (!(o & o - 1))
ej[i] = (int *) realloc(ej[i], (o << 1) * sizeof *ej[i]);
ej[i][o] = j;
}
void detach(int i, int j) {
for (int o = 0; o < eo[i]; o++)
if (ej[i][o] == j) {
swap(ej[i][o], ej[i][--eo[i]]);
return;
}
}
void dfs0(int p, int i) {
detach(i, p);
for (int k = 0; k <= k_; k++)
dp[i][k][0] = 0, dp[i][k][1] = k ? bb[i] : -INF;
for (int o = 0; o < eo[i]; o++) {
int j = ej[i][o];
dfs0(i, j);
for (int k = 0; k <= k_; k++) {
dp[i][k][0] = max(dp[i][k][0], max(dp[j][k][0], dp[j][k][1]));
if (k < k_)
dp[i][k + 1][1] = max(dp[i][k + 1][1], max(dp[j][k][0], dp[j][k][1]) + bb[i] - aa[j]);
}
}
}
void dfs1(int p, int i) {
if (p != -1)
for (int k = 0; k <= k_; k++)
dq[i][k] = max(0LL, k ? bb[i] - aa[p] : -INF);
for (int o = 0; o < eo[i]; o++) {
int j = ej[i][o];
dfs1(i, j);
if (p != -1)
for (int k = 0; k <= k_; k++)
dq[i][k] = max(dq[i][k], max(dq[j][k], k ? dq[j][k - 1] + bb[i] - aa[p] : -INF));
}
}
int main() {
ios_base::sync_with_stdio(false), cin.tie(NULL);
int n; cin >> n >> k_;
for (int i = 0; i < n; i++)
cin >> aa[i];
for (int h = 0; h < n - 1; h++) {
int i, j; cin >> i >> j, i--, j--;
append(i, j), bb[i] += aa[j];
append(j, i), bb[j] += aa[i];
}
dfs0(-1, 0);
dfs1(-1, 0);
long long ans = 0;
for (int i = 0; i < n; i++) {
for (int f = 0; f < 2; f++)
ans = max(ans, dp[i][k_][f]);
for (int o = 0; o < eo[i]; o++) {
int j = ej[i][o];
ans = max(ans, max(dq[j][k_], k_ ? dq[j][k_ - 1] + bb[i] : -INF));
}
}
for (int i = 0; i < n; i++) {
for (int k = 0; k <= k_; k++)
xx[k] = -INF;
for (int o = 0; o < eo[i]; o++) {
int j = ej[i][o];
for (int k = 0; k <= k_; k++) {
ans = max(ans, max(dp[j][k][0], dp[j][k][1]) + xx[k_ - k]);
if (k < k_)
ans = max(ans, max(dp[j][k][0], dp[j][k][1]) + bb[i] - aa[j] + xx[k_ - k - 1]);
}
for (int k = 0; k <= k_; k++)
xx[k] = max(xx[k], dq[j][k]);
}
for (int k = 0; k <= k_; k++)
xx[k] = -INF;
for (int o = eo[i] - 1; o >= 0; o--) {
int j = ej[i][o];
for (int k = 0; k <= k_; k++) {
ans = max(ans, max(dp[j][k][0], dp[j][k][1]) + xx[k_ - k]);
if (k < k_)
ans = max(ans, max(dp[j][k][0], dp[j][k][1]) + bb[i] - aa[j] + xx[k_ - k - 1]);
}
for (int k = 0; k <= k_; k++)
xx[k] = max(xx[k], dq[j][k]);
}
}
cout << ans << '\n';
return 0;
}
# | 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... |