Submission #860154

#TimeUsernameProblemLanguageResultExecution timeMemory
860154ThegeekKnight16Chase (CEOI17_chase)C++17
40 / 100
4051 ms494572 KiB
#include <bits/stdc++.h>
using namespace std;
#define int long long
const int MAXN = 1e5 + 10;
const int MAXV = 110;
array<vector<array<int, 3>>, MAXN> dp;
array<vector<int>, MAXN> grafo;
array<int, MAXN> pom, pomviz;
int N, V;
int resp = 0;

void dfs(int v, int p)
{
    dp[v].resize(V+1, {0, 0, 0});
    // for (int i = 0; i <= V; i++) dp[v][i].emplace(0, v);
    pomviz[v] = 0;
    for (auto viz : grafo[v])
    {
        if (viz == p) continue;
        pomviz[v] += pom[viz];
    }

    for (auto viz : grafo[v])
    {
        if (viz == p) continue;
        dfs(viz, v);
        for (int k = 1; k <= V; k++)
        {
            int prox = max(dp[viz][k][0], dp[viz][k-1][0] + pomviz[v]);
            auto &atual = dp[v][k];
            if (prox > atual[0])
            {
                swap(atual[2], atual[1]);
                swap(atual[1], atual[0]);
                atual[0] = prox;
            }
            else if (prox > atual[1])
            {
                swap(atual[2], atual[1]);
                atual[1] = prox;
            }
            else if (prox > atual[2])
            {
                atual[2] = prox;
            }
        }
    }
}

void girarArvore(int v, int p)
{
    // cerr << "v: " << v << '\n';
    // cerr << '\t' << "dp: ";
    // for (auto x : dp[v][V]) cerr << get<0>(x) << " ";
    // cerr << '\n';
    // cerr << '\t' << "pomviz: " << pomviz[v] << '\n';
    resp = max(resp, dp[v][V][0]);

    for (auto viz : grafo[v])
    {
        if (viz == p) continue;

        //Salva
        auto antv = dp[v]; int antpom = pomviz[v];
        auto antviz = dp[viz]; int antpomviz = pomviz[viz];

        //Gira
        pomviz[v] -= pom[viz]; pomviz[viz] += pom[v];
        for (int k = 1; k <= V; k++)
        {
            auto &atual = dp[v][k];
            atual = {0, 0, 0};
            for (auto viz2 : grafo[v])
            {
                if (viz2 == viz) continue;
                int prox = max(dp[viz2][k][0], dp[viz2][k-1][0] + pomviz[v]);
                if (prox > atual[0])
                {
                    swap(atual[2], atual[1]);
                    swap(atual[1], atual[0]);
                    atual[0] = prox;
                }
                else if (prox > atual[1])
                {
                    swap(atual[2], atual[1]);
                    atual[1] = prox;
                }
                else if (prox > atual[2])
                {
                    atual[2] = prox;
                }
            }
        }
        for (int k = 1; k <= V; k++)
        {
            auto &atual = dp[viz][k];
            atual = {0, 0, 0};
            for (auto viz2 : grafo[viz])
            {
                int prox = max(dp[viz2][k][0], dp[viz2][k-1][0] + pomviz[viz]);
                if (prox > atual[0])
                {
                    swap(atual[2], atual[1]);
                    swap(atual[1], atual[0]);
                    atual[0] = prox;
                }
                else if (prox > atual[1])
                {
                    swap(atual[2], atual[1]);
                    atual[1] = prox;
                }
                else if (prox > atual[2])
                {
                    atual[2] = prox;
                }
            }
        }

        girarArvore(viz, v);

        //Volta
        dp[v] = antv; dp[viz] = antviz;
        pomviz[v] = antpom; pomviz[viz] = antpomviz;
    }
}

int32_t main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cin >> N >> V;
    for (int i = 1; i <= N; i++) cin >> pom[i];
    for (int i = 1; i < N; i++)
    {
        int X, Y;
        cin >> X >> Y;
        grafo[X].push_back(Y);
        grafo[Y].push_back(X);
    }
    dfs(1, 0);
    girarArvore(1, 1);
    cout << resp << '\n';
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...