Submission #1206646

#	Time	Username	Problem	Language	Result	Execution time	Memory
1206646		ocean	Team Coding (EGOI24_teamcoding)	C++17	42 / 100	509 ms	1114112 KiB

Main

#include <iostream>
#include <vector>
#include <algorithm>
#include <map> // Not strictly needed with array-based frequency counts if K_langs and N_nodes are within limits for arrays

using namespace std;

const long long INF = 4e18; // P can be N, S can be N. Sum can be 2N. Using a larger INF.

vector<int> adj[100005];
int emp_lang[100005]; // Renamed to avoid conflict
int level[100005];
int subtree_sz[100005];
int n_nodes, k_langs;

long long max_p_global = 0;
long long min_s_global = INF;

vector<vector<int>> global_C; 
long long current_P_val[100005]; 
long long current_I_val[100005]; 
int freq_level[100005]; 


void dfs_levels(int u, int p, int l) {
    level[u] = l;
    for (int v : adj[u]) {
        if (v == p) continue;
        dfs_levels(v, u, l + 1);
    }
}

void dfs_size(int u, int p) {
    subtree_sz[u] = 1;
    for (int v : adj[u]) {
        if (v == p) continue;
        dfs_size(v, u);
        subtree_sz[u] += subtree_sz[v];
    }
}

void update_sack_data(int node_x, int val_sign) {
    int lvl_x = level[node_x];
    int lang_x = emp_lang[node_x];

    long long old_freq_at_lvl_x = freq_level[lvl_x];
    freq_level[lvl_x] += val_sign;
    long long new_freq_at_lvl_x = freq_level[lvl_x];

    for (int k_test = 0; k_test < k_langs; ++k_test) {
        long long term_change = 0;
        // Subtract contribution of old_freq_at_lvl_x
        if (old_freq_at_lvl_x > 0) {
             term_change -= min(old_freq_at_lvl_x, (long long)global_C[k_test][lvl_x]);
        }
        // Add contribution of new_freq_at_lvl_x
        if (new_freq_at_lvl_x > 0) {
             term_change += min(new_freq_at_lvl_x, (long long)global_C[k_test][lvl_x]);
        }
        if (term_change != 0) { // Optimization: only update if there's a change
            current_P_val[k_test] += term_change;
        }
    }
    current_I_val[lang_x] += val_sign;
}

void dfs_process_subtree(int u, int p, int val_sign) {
    update_sack_data(u, val_sign);
    for (int v : adj[u]) {
        if (v == p) continue;
        dfs_process_subtree(v, u, val_sign);
    }
}

void dfs_sack(int u, int p, bool keep) {
    int max_child_sz = -1, big_child = -1;
    for (int v : adj[u]) {
        if (v == p) continue;
        if (subtree_sz[v] > max_child_sz) {
            max_child_sz = subtree_sz[v];
            big_child = v;
        }
    }

    for (int v : adj[u]) {
        if (v == p || v == big_child) continue;
        dfs_sack(v, u, false); 
    }

    if (big_child != -1) {
        dfs_sack(big_child, u, true); 
    }

    update_sack_data(u, 1);

    for (int v : adj[u]) {
        if (v == p || v == big_child) continue;
        dfs_process_subtree(v, u, 1);
    }

    int k_u = emp_lang[u];
    long long cand_P = current_P_val[k_u];
    long long cand_S = cand_P - current_I_val[k_u];

    if (cand_P > max_p_global) {
        max_p_global = cand_P;
        min_s_global = cand_S;
    } else if (cand_P == max_p_global) {
        min_s_global = min(min_s_global, cand_S);
    }
    
    if (!keep) {
        dfs_process_subtree(u, p, -1); 
    }
}


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

    cin >> n_nodes >> k_langs;

    for (int i = 0; i < n_nodes; ++i) {
        cin >> emp_lang[i];
    }
    
    // Build directed adjacency list for the tree
    vector<vector<int>> adj_build(n_nodes);
    for (int i = 1; i < n_nodes; ++i) { 
        int b_val;
        cin >> b_val; 
        adj_build[b_val].push_back(i);
    }
    for(int i=0; i<n_nodes; ++i) adj[i] = adj_build[i];


    dfs_levels(0, -1, 0);

    global_C.resize(k_langs, vector<int>(n_nodes + 1, 0)); 
    for (int i = 0; i < n_nodes; ++i) {
        if (level[i] < n_nodes) { 
             global_C[emp_lang[i]][level[i]]++;
        }
    }
    
    dfs_size(0, -1); 

    for(int i=0; i<k_langs; ++i) {
        current_P_val[i] = 0;
        current_I_val[i] = 0;
    }
    for(int i=0; i<n_nodes; ++i) { 
        freq_level[i] = 0;
    }
    
    // Initialize min_s_global for the case n_nodes = 1
    if (n_nodes == 1) {
         max_p_global = 1;
         min_s_global = 0;
    }


    dfs_sack(0, -1, false); 

    cout << max_p_global << " " << min_s_global << endl;

    return 0;
}

• Subtask 1: 0 / 12
• Subtask 2: 19 / 19
• Subtask 3: 0 / 27
• Subtask 4: 23 / 23
• Subtask 5: 0 / 19