Submission #519386

#TimeUsernameProblemLanguageResultExecution timeMemory
519386hoanghq2004Chase (CEOI17_chase)C++14
100 / 100
317 ms268384 KiB
#include <bits/stdc++.h>
#pragma GCC optimization ("O3")
#pragma GCC optimization ("unroll-loops")
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace __gnu_pbds;
using namespace std;

template <typename T>
using ordered_set = tree <T, null_type, less <T>, rb_tree_tag, tree_order_statistics_node_update>;

const int N = 1e5 + 10;

int n, m;
long long in[N][110];
long long out[N][110];
int c[N];
long long tot[N];
vector <int> g[N];
long long prf[N][110], suf[N][110];
int child[N];

void down(int u, int p) {
    for (auto v: g[u]) {
        if (v == p) continue;
        down(v, u);
        for (int i = 1; i <= m; ++i)
            in[u][i] = max({in[u][i], in[v][i], in[v][i - 1] + tot[v] - c[u]});
    }
}

void up(int u, int p) {
    int num = 0;
    for (auto v: g[u]) {
        if (v == p) continue;
        child[++num] = v;
    }
    for (int i = 0; i <= m; ++i) prf[0][i] = suf[num + 1][i];
    for (int i = 1; i <= num; ++i) {
        int v = child[i];
        for (int j = 1; j <= m; ++j) {
            prf[i][j] = max({prf[i - 1][j], in[v][j], in[v][j - 1] + tot[v] - c[u]});
        }
    }
    for (int i = num; i >= 1; --i) {
        int v = child[i];
        for (int j = 1; j <= m; ++j) {
            suf[i][j] = max({suf[i + 1][j], in[v][j], in[v][j - 1] + tot[v] - c[u]});
        }
    }
    for (int i = 1; i <= num; ++i) {
        int v = child[i];
        for (int j = 1; j <= m; ++j) {
            out[v][j] = max(out[u][j], out[u][j - 1] + tot[u] - c[v]);
            out[v][j] = max({out[v][j], prf[i - 1][j], suf[i + 1][j]});
            out[v][j] = max({out[v][j], prf[i - 1][j - 1] + tot[u] - c[v], suf[i + 1][j - 1] + tot[u] - c[v]});
        }
    }
    for (auto v: g[u]) {
        if (v == p) continue;
        up(v, u);
    }
}

int main() {
    ios :: sync_with_stdio(0); cin.tie(0);
    cin >> n >> m;
    for (int i = 1; i <= n; ++i) cin >> c[i];
    for (int i = 1; i < n; ++i) {
        int u, v;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    for (int u = 1; u <= n; ++u)
        for (auto v: g[u])
            tot[u] += c[v];
    down(1, 0);
    up(1, 0);
    long long ans = 0;
    for (int u = 1; u <= n; ++u) {
        ans = max({ans, in[u][m], out[u][m]});
        for (int i = 0; i < m; ++i) {
            ans = max({ans, in[u][i] + tot[u], out[u][i] + tot[u]});
        }
    }
    cout << ans;
}

Compilation message (stderr)

chase.cpp:2: warning: ignoring '#pragma GCC optimization' [-Wunknown-pragmas]
    2 | #pragma GCC optimization ("O3")
      | 
chase.cpp:3: warning: ignoring '#pragma GCC optimization' [-Wunknown-pragmas]
    3 | #pragma GCC optimization ("unroll-loops")
      |
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...