Submission #593469

#	Time	Username	Problem	Language	Result	Execution time	Memory
593469		alextodoran	Chase (CEOI17_chase)	C++17	100 / 100	462 ms	180060 KiB

chase

/**
 ____ ____ ____ ____ ____
||a |||t |||o |||d |||o ||
||__|||__|||__|||__|||__||
|/__\|/__\|/__\|/__\|/__\|

**/

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;

const int N_MAX = (int) 1e5;
const int K_MAX = 100;

int N, K;
int P[N_MAX + 2];
vector <int> adj[N_MAX + 2];

ll down[N_MAX + 2][K_MAX + 2];
ll up[N_MAX + 2][K_MAX + 2];
ll full[N_MAX + 2];

ll maxDown[K_MAX + 2];

void dfs (int u, int parent = 0) {
    vector <int> sons;
    ll sumSons = 0;
    for (int v : adj[u]) {
        if (v != parent) {
            sons.push_back(v);
            dfs(v, u);
            sumSons += P[v];
        }
    }
    down[u][0] = 0;
    down[u][1] = sumSons;
    up[u][0] = 0;
    up[u][1] = sumSons + P[parent];
    for (int v : sons) {
        for (int k = 0; k <= K; k++) {
            down[u][k] = max(down[u][k], down[v][k]);
            up[u][k] = max(up[u][k], up[v][k]);
            if (k > 0) {
                down[u][k] = max(down[u][k], down[v][k - 1] + sumSons);
                up[u][k] = max(up[u][k], up[v][k - 1] + sumSons - P[v] + P[parent]);
            }
        }
        full[u] = max(full[u], full[v]);
    }
    for (int k = 1; k <= K; k++) {
        down[u][k] = max(down[u][k], down[u][k - 1]);
        up[u][k] = max(up[u][k], up[u][k - 1]);
    }
    fill(maxDown, maxDown + K + 1, 0);
    for (int v : sons) {
        for (int k = 0; k <= K; k++) {
            full[u] = max(full[u], up[v][k] + maxDown[K - k]);
            if (k < K) {
                full[u] = max(full[u], up[v][k] + maxDown[K - k - 1] + sumSons - P[v] + P[parent]);
            }
            full[u] = max(full[u], down[v][k]);
            if (k > 0) {
                full[u] = max(full[u], down[v][k - 1] + sumSons + P[parent]);
            }
        }
        for (int k = 0; k <= K; k++) {
            maxDown[k] = max(maxDown[k], down[v][k]);
        }
    }
    reverse(sons.begin(), sons.end());
    fill(maxDown, maxDown + K + 1, 0);
    for (int v : sons) {
        for (int k = 0; k <= K; k++) {
            full[u] = max(full[u], up[v][k] + maxDown[K - k]);
            if (k < K) {
                full[u] = max(full[u], up[v][k] + maxDown[K - k - 1] + sumSons - P[v] + P[parent]);
            }
        }
        for (int k = 0; k <= K; k++) {
            maxDown[k] = max(maxDown[k], down[v][k]);
        }
    }
}

int main () {
    ios_base::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);

    cin >> N >> K;
    for (int u = 1; u <= N; u++) {
        cin >> P[u];
    }
    for (int i = 1; i <= N - 1; i++) {
        int u, v;
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    dfs(1);
    cout << full[1] << "\n";

    return 0;
}

• Subtask 1: 20 / 20
• Subtask 2: 20 / 20
• Subtask 3: 30 / 30
• Subtask 4: 30 / 30