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 std::vector;
using std::array;
using std::pair;
using std::tuple;
using i64 = std::int64_t;
template <class F> struct RecLambda : private F {
explicit RecLambda(F&& f) : F(std::forward<F>(f)) {}
template <class... Args> decltype(auto) operator()(Args&&... args) const {
return F::operator()(*this, std::forward<Args>(args)...);
}
};
template <class T> void setmax(T& x, const T& y) {
if (x < y) x = y;
}
int main() {
std::ios_base::sync_with_stdio(false);
std::cin.tie(nullptr);
int N, K;
std::cin >> N >> K;
if (K == 0) {
std::cout << 0 << '\n';
return 0;
}
vector<int> W(N);
for (auto& x : W) {
std::cin >> x;
}
vector<vector<int>> graph(N);
for (int i = 0; i < N - 1; ++i) {
int a, b;
std::cin >> a >> b;
a -= 1, b -= 1;
graph[a].push_back(b);
graph[b].push_back(a);
}
vector<i64> C(N);
for (int i = 0; i < N; ++i) {
for (const int j : graph[i]) {
C[i] += W[j];
}
}
i64 ans = 0;
RecLambda([&](auto&& dfs, const int u, const int p) -> pair<vector<i64>, vector<i64>> {
vector<i64> begin(K + 1), end(K + 1);
begin[1] = end[1] = C[u];
for (const int v : graph[u]) {
if (v == p) continue;
const auto [v_begin, v_end] = dfs(v, u);
for (int i = 1; i < K; ++i) {
setmax(ans, end[i] + v_begin[K - i] - W[u]);
setmax(ans, end[i] + v_begin[K - i + 1] - C[v]);
setmax(ans, begin[i] + v_end[K - i] - W[v]);
setmax(ans, begin[i + 1] + v_end[K - i] - C[u]);
}
for (int i = 0; i < K; ++i) {
setmax(begin[i + 1], C[u] + v_begin[i + 1] - C[v]);
setmax(begin[i + 1], C[u] + v_begin[i] - W[u]);
setmax(end[i], v_end[i]);
setmax(end[i + 1], v_end[i] + C[u] - W[v]);
}
setmax(end[K], v_end[K]);
}
for (int i = 0; i < K; ++i) {
setmax(begin[i + 1], begin[i]);
setmax(end[i + 1], end[i]);
}
setmax(ans, begin[K]);
setmax(ans, end[K]);
return std::pair(std::move(begin), std::move(end));
})(0, -1);
std::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... |