Submission #960964

#	Time	Username	Problem	Language	Result	Execution time	Memory
960964		Ghetto	Chase (CEOI17_chase)	C++17	100 / 100	3234 ms	120336 KiB

chase

#include <bits/stdc++.h>
using namespace std;
using pii = pair<int, int>;
using lint = long long;
const int MAX_N = 1e5 + 5, MAX_K = 1e2 + 5;

int n, k;
lint val[MAX_N];
vector<int> adj[MAX_N], adj_ind[MAX_N]; // adj_ind[u][i] = ind of adj list u is for adj[u][i]

int par[MAX_N];
vector<int> children[MAX_N];
void init1() {
    for (int i = 1; i <= n; i++) children[i].push_back(0);
}
int child_ind[MAX_N];

int n_inds, ind[MAX_N], rev_ind[MAX_N];
void dfs(int u) {
    ind[u] = ++n_inds;
    rev_ind[ind[u]] = u;
    for (int v : adj[u]) {
        if (ind[v]) continue;
        par[v] = u;
        children[u].push_back(v);
        child_ind[v] = children[u].size() - 1;
        dfs(v);
    }
}

lint val_sum[MAX_N];

void precomp() {
    init1();

    dfs(1);

    for (int i = 1; i <= n; i++)
        for (int j : adj[i]) val_sum[i] += val[j];
}

lint dp1[MAX_N][MAX_K];
vector<lint> dp2[MAX_N][2];
lint max_dp2[MAX_N][2];
void init2() {
    for (int i = 1; i <= n; i++)
        for (int c = 0; c <= 1; c++)
            dp2[i][c].resize(adj[i].size() + 2);
}

lint pref_max[MAX_N], suff_max[MAX_N];

int main() {
    // freopen("chase.in", "r", stdin);

    cin >> n >> k;
    for (int i = 1; i <= n; i++) cin >> val[i];
    for (int i = 1; i < n; i++) {
        int u, v; cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    precomp();

    for (int c = 1; c <= k; c++) {
        for (int i = n; i >= 1; i--) {
            int u = rev_ind[i];

            lint leave = 0;
            for (int v : children[u])
                if (v != 0) leave = max(leave, dp1[v][c]);

            lint take = 0; 
            for (int v : children[u])
                if (v != 0) take = max(take, dp1[v][c - 1]);
            take += val_sum[u] - val[par[u]];

            dp1[u][c] = max(take, leave);
        }
    }

    init2();
    lint ans = 0;
    for (int c = 0; c <= k; c++) {
        int opp_c = k - c;
        int p = c % 2, opp_p = (c + 1) % 2;
        for (int i = n; i >= 1; i--) {
            int u = rev_ind[i];

            max_dp2[u][p] = 0;
            for (int j = 0; j < children[u].size(); j++) {
                pref_max[j] = max(((j == 0) ? 0 : pref_max[j - 1]), dp1[children[u][j]][opp_c]);
                if (c == 0) continue;
                int v = children[u][j];

                lint leave = (v == 0) ? 0 : max_dp2[v][p];
                lint take = val_sum[u] - val[v] + ((v == 0) ? 0 : max_dp2[v][opp_p]);
                dp2[u][p][j] = max(take, leave);
                max_dp2[u][p] = max(max_dp2[u][p], dp2[u][p][j]);
            }

            for (int j = children[u].size() - 1; j >= 0; j--) 
                suff_max[j] = max((j == children[u].size() - 1) ? 0 : suff_max[j + 1], dp1[children[u][j]][opp_c]);
            
            for (int j = 0; j < children[u].size(); j++) {               
                lint tot_max = (j == 0) ? 0 : pref_max[j - 1];
                tot_max = max(tot_max, ((j == children[u].size() - 1) ? 0 : suff_max[j + 1]));

                ans = max(ans, dp2[u][p][j] + tot_max);
            }
        }
    }
    cout << ans << '\n';
    
}

Compilation message (stderr)

chase.cpp: In function 'int main()':
chase.cpp:92:31: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   92 |             for (int j = 0; j < children[u].size(); j++) {
      |                             ~~^~~~~~~~~~~~~~~~~~~~
chase.cpp:104:38: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  104 |                 suff_max[j] = max((j == children[u].size() - 1) ? 0 : suff_max[j + 1], dp1[children[u][j]][opp_c]);
      |                                    ~~^~~~~~~~~~~~~~~~~~~~~~~~~
chase.cpp:106:31: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  106 |             for (int j = 0; j < children[u].size(); j++) {
      |                             ~~^~~~~~~~~~~~~~~~~~~~
chase.cpp:108:44: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  108 |                 tot_max = max(tot_max, ((j == children[u].size() - 1) ? 0 : suff_max[j + 1]));
      |                                          ~~^~~~~~~~~~~~~~~~~~~~~~~~~

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