This submission is migrated from previous version of oj.uz, which used different machine for grading. This submission may have different result if resubmitted.
// Success consists of going from failure to failure without loss of enthusiasm
#include <bits/stdc++.h>
using namespace std;
#define nl '\n'
#define pb push_back
#define sz(x) int(x.size())
using ll = int64_t;
template<class T> using V = vector<T>;
using vl = V<ll>;
using vi = V<int>;
const int nax = 1e5+5;
const int kax = 105;
const ll INFL = 1e18;
ll up[nax][kax], dwn[nax][kax];
int main() {
cin.tie(0)->sync_with_stdio(0);
int N, K; cin >> N >> K;
vl A(N); for(auto& x : A) cin >> x;
vl P(N);
V<vi> adj(N); for(int i = 0; i < N - 1; i++) {
int u, v; cin >> u >> v; --u, --v;
adj[u].pb(v);
adj[v].pb(u);
P[u] += A[v];
P[v] += A[u];
}
ll ans = 0;
memset(up, 0, sizeof up);
memset(dwn, 0, sizeof dwn);
if (K > 0) for(int i = 0; i < N; i++) up[i][1] = P[i];
// up includes u, while dwn does not
V<vl> bst(K + 1, vl(4, -1));
function<void(int, int)> gen = [&](int u, int p) {
// cout << u << " " << p << endl;
auto UPV = [&] (int v, int k) {
if (v == -1) return ll(min(1, k) * P[u]);
return max(up[v][k], up[v][k - 1] + P[u] - A[v]);
};
auto DWNV = [&](int v, int k) {
if (v == -1) return ll(0);
return max(dwn[v][k], dwn[v][k - 1] + P[v] - A[u]);
};
vi chd;
for(auto& v : adj[u]) if (v != p) {
gen(v, u); chd.pb(v);
}
for(auto& v : chd) {
for(int k = 1; k <= K; k++) {
up[u][k] = max(up[u][k], UPV(v, k));
dwn[u][k] = max(dwn[u][k], DWNV(v, k));
if (UPV(bst[k][1], k) < UPV(v, k)) {
bst[k][1] = v;
if (UPV(bst[k][0], k) < UPV(bst[k][1], k)) swap(bst[k][1], bst[k][0]);
}
if (DWNV(bst[k][3], k) < DWNV(v, k)) {
bst[k][3] = v;
if (DWNV(bst[k][2], k) < DWNV(bst[k][3], k)) swap(bst[k][3], bst[k][2]);
}
}
}
for(int k = 0; k <= K; k++) {
ans = max({ans, up[u][k], dwn[u][k]});
int nk = K - k;
// k -> up, nk -> dwn
for(int x = 0; x < 2; x++) for(int y = 2; y < 4; y++) {
if (bst[k][x] == bst[nk][y]) continue;
ans = max(ans, UPV(bst[k][x], k) + DWNV(bst[nk][y], nk));
}
bst[k][0] = bst[k][1] = bst[nk][2] = bst[nk][3] = -1;
}
};
gen(0, -1);
cout << ans << nl;
exit(0-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... |