Submission #1206648

#	Time	Username	Problem	Language	Result	Execution time	Memory
1206648		ocean	Team Coding (EGOI24_teamcoding)	C++17	54 / 100	4086 ms	40260 KiB

Main

#include <iostream>
#include <vector>
#include <algorithm>
#include <map>

using namespace std;

const long long INF = 4e18; // Sufficiently large for min_s_global.

vector<int> adj[100005];
int emp_lang[100005];
int level[100005];
int subtree_sz[100005];
int n_nodes, k_langs;

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

// Sparse representation of C[k][l]
// relevant_Cs_at_level[l] stores a list of pairs {lang_k, count}, 
// where count = C[lang_k][l] and is > 0.
vector<vector<pair<int, int>>> relevant_Cs_at_level;

// Global arrays for Sack states
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];

    // Update current_P_val[k_test] for relevant k_test based on change at lvl_x
    // Only languages k_test for which C[k_test][lvl_x] > 0 are affected.
    if (lvl_x < (int)relevant_Cs_at_level.size()) { // Boundary check for lvl_x
        for (auto const& pair_kc : relevant_Cs_at_level[lvl_x]) {
            int k_test = pair_kc.first;
            long long C_val_k_test_lvl_x = pair_kc.second; 

            long long term_change = 0;
            if (old_freq_at_lvl_x > 0) { 
                term_change -= min(old_freq_at_lvl_x, C_val_k_test_lvl_x);
            }
            if (new_freq_at_lvl_x > 0) { 
                term_change += min(new_freq_at_lvl_x, C_val_k_test_lvl_x);
            }
            // No specific check for term_change != 0 needed, X += 0 is fine.
            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); // Add u itself

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

    // Query for u as team lead
    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) {
        // Remove S_u (u and all its children including big child if it was processed under u)
        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);

    // Precompute sparse C[k][l]
    relevant_Cs_at_level.resize(n_nodes); // Max level is n_nodes-1
    vector<map<int, int>> temp_counts_for_Cs(n_nodes); 
    for (int i = 0; i < n_nodes; ++i) {
        if (level[i] < n_nodes) { // Should always be true
             temp_counts_for_Cs[level[i]][emp_lang[i]]++;
        }
    }
    for (int l = 0; l < n_nodes; ++l) {
        if (!temp_counts_for_Cs[l].empty()) {
            relevant_Cs_at_level[l].reserve(temp_counts_for_Cs[l].size()); 
            for (auto const& pair_kl : temp_counts_for_Cs[l]) { 
                relevant_Cs_at_level[l].push_back({pair_kl.first, pair_kl.second});
            }
        }
    }
    
    // Initialize Sack global variables (static arrays are zero-initialized if global)
    // Explicitly zeroing for clarity, though not strictly needed for global static arrays.
    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;
    }
    
    dfs_size(0, -1); 
    
    dfs_sack(0, -1, false); 

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

    return 0;
}

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