Submission #523655

#TimeUsernameProblemLanguageResultExecution timeMemory
523655KoDChase (CEOI17_chase)C++17
100 / 100
184 ms102276 KiB
#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 timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...