Submission #1206657

#	Time	Username	Problem	Language	Result	Execution time	Memory
1206657		ocean	Team Coding (EGOI24_teamcoding)	C++17	54 / 100	4094 ms	40264 KiB

Main

// Standard library includes
#include <iostream> // For input/output operations (cin, cout)
#include <vector>   // For dynamic arrays (std::vector)
#include <algorithm>// For std::min, std::max, etc.
#include <map>      // For std::map (used in precomputing C[k][l])

// Using standard namespace for convenience, to avoid writing std:: repeatedly
using namespace std; 

// INF is a large value used to initialize min_s_global (minimum switches).
// Team size P and switches S are counts of people, at most N.
// current_P_val can accumulate sums up to N*N in theory (N levels, each term min(N,N)), so long long is needed.
// Final max_p_global and min_s_global will be at most N. Using long long for them for consistency.
const long long INF = 4e18; 

// --- Global Variables ---
// These variables are accessible throughout the program.

// Adjacency list to represent the company hierarchy as a tree.
// adj[i] will store a vector of employee IDs who directly report to employee i.
vector<int> adj[100005]; 
// emp_lang[i] stores the preferred programming language of employee i.
int emp_lang[100005];    
// level[i] stores the depth of employee i in the tree (CEO/root is at level 0).
int level[100005];       
// subtree_sz[i] stores the number of employees in the subtree rooted at employee i (including i).
// This is used by the Sack algorithm to find the "big child".
int subtree_sz[100005];  

// n_nodes: total number of employees.
// k_langs: total number of different programming languages.
int n_nodes, k_langs; 

// max_p_global: stores the maximum team size found so far across all possible team leads.
long long max_p_global = 0;    
// min_s_global: stores the minimum number of switches required to achieve max_p_global.
long long min_s_global = INF;  

// Sparse representation of C[k][l]:
// C[k][l] would be the total count of employees in the *entire company*
// who prefer language 'k' and are at level 'l'.
// Instead of a dense KxN matrix (which would be too large for memory),
// relevant_Cs_at_level[l] stores a list of pairs {language_idx, count}.
// This list only includes languages that actually have 'count' > 0 employees at level 'l'.
// The total number of pairs stored across all levels is at most n_nodes.
vector<vector<pair<int, int>>> relevant_Cs_at_level;

// Global arrays for the Sack algorithm's state:
// As the Sack algorithm processes a subtree (rooted at some 'u'), these arrays
// are updated to reflect aggregated information about the nodes currently
// considered "in the Sack's bag" (i.e., nodes from S_u).

// freq_level[l]: stores the count of nodes currently in the Sack bag that are at global tree level 'l'.
// Represents |{x in S_u | level[x]==l}| for the current S_u.
int freq_level[100005];          
// current_I_val[lang_idx]: stores the count of nodes currently in the Sack bag
// that prefer language 'lang_idx'.
// Represents sum_{x in S_u, emp_lang[x]==lang_idx} 1. This is used to find initial_members_u.
long long current_I_val[100005]; 
// current_P_val[k_test]: for a given language k_test, this stores the sum over all levels l of:
// min(freq_level[l], C[k_test][l]).
// This represents the total team size if k_test was the project language, based on the current Sack bag.
long long current_P_val[100005]; 


// --- Helper DFS Functions (executed once at the beginning) ---

// Computes the level (depth) of each employee in the tree.
// CEO (node 0) is at level 0.
// 'u': current employee being visited.
// 'p': parent of 'u' in the DFS traversal (to avoid going back up).
// 'l': current level to assign to 'u'.
void dfs_levels(int u, int p, int l) {
    level[u] = l; // Assign current level 'l' to employee 'u'
    for (int v : adj[u]) { // For each direct report 'v' of 'u'
        if (v == p) continue; // Skip the parent node
        dfs_levels(v, u, l + 1); // Recursively call for child 'v' at the next level (l+1)
    }
}

// Computes the size of the subtree rooted at each employee 'u'.
// Subtree size includes 'u' itself.
// This is needed by the Sack algorithm to identify the "big child" for optimization.
// 'u': current employee being visited.
// 'p': parent of 'u'.
void dfs_size(int u, int p) {
    subtree_sz[u] = 1; // Employee 'u' itself contributes 1 to its subtree size
    for (int v : adj[u]) { // For each direct report 'v' of 'u'
        if (v == p) continue; // Skip parent
        dfs_size(v, u); // Recursively compute subtree size for child 'v'
        subtree_sz[u] += subtree_sz[v]; // Add child's subtree size to 'u's subtree size
    }
}


// --- Sack Algorithm Core Functions ---

// This function updates the global Sack state arrays (freq_level, current_P_val, current_I_val)
// when 'node_x' is either added to or removed from the Sack's consideration.
// 'node_x': the employee node being processed.
// 'val_sign': +1 if adding node_x, -1 if removing node_x.
void update_sack_data(int node_x, int val_sign) {
    int lvl_x = level[node_x];    // Get the level of node_x
    int lang_x = emp_lang[node_x]; // Get the language of node_x

    // 1. Update freq_level and get old/new values
    // freq_level[lvl_x] tracks how many nodes currently in the Sack bag are at this level.
    long long old_freq_at_lvl_x = freq_level[lvl_x]; // Count before this node_x is processed
    freq_level[lvl_x] += val_sign;                   // Add/remove node_x's presence at this level
    long long new_freq_at_lvl_x = freq_level[lvl_x]; // Count after node_x is processed

    // 2. Update current_P_val[k_test] for all relevant languages k_test.
    // current_P_val[k_test] = sum over all levels l of: min(freq_level[l], C[k_test][l]).
    // When freq_level[lvl_x] changes, only the term for l=lvl_x in this sum changes.
    // The change is: min(new_freq_at_lvl_x, C[k_test][lvl_x]) - min(old_freq_at_lvl_x, C[k_test][lvl_x]).
    // We only need to do this for languages k_test that actually have employees
    // at lvl_x in the whole company (i.e., C[k_test][lvl_x] > 0).
    // This is where `relevant_Cs_at_level` (our sparse C matrix) is used.
    if (lvl_x < (int)relevant_Cs_at_level.size()) { // Boundary check for level index
        // Iterate through languages k_test that have employees at lvl_x
        for (auto const& pair_kc : relevant_Cs_at_level[lvl_x]) {
            int k_test = pair_kc.first;                 // The language k_test
            long long C_val_k_test_lvl_x = pair_kc.second; // C[k_test][lvl_x] (total employees in company with lang k_test at lvl_x)

            long long term_change = 0;
            // Subtract the contribution of lvl_x to P_val[k_test] based on old_freq_at_lvl_x
            if (old_freq_at_lvl_x > 0) { // If there were nodes at this level in the bag before
                term_change -= min(old_freq_at_lvl_x, C_val_k_test_lvl_x);
            }
            // Add the contribution of lvl_x to P_val[k_test] based on new_freq_at_lvl_x
            if (new_freq_at_lvl_x > 0) { // If there are nodes at this level in the bag now
                term_change += min(new_freq_at_lvl_x, C_val_k_test_lvl_x);
            }
            // Apply the net change to current_P_val for this language k_test
            current_P_val[k_test] += term_change;
        }
    }
    
    // 3. Update current_I_val for the language of node_x.
    // current_I_val[lang_x] tracks the count of nodes in the Sack bag preferring lang_x.
    // This directly changes by val_sign (+1 for add, -1 for remove).
    current_I_val[lang_x] += val_sign;
}

// Helper DFS to apply update_sack_data to all nodes in a given subtree.
// This is used by the Sack algorithm to add contributions of small children's subtrees
// or to clear data from a subtree when it's no longer needed.
// 'u': root of the subtree to process.
// 'p': parent of 'u'.
// 'val_sign': +1 to add subtree nodes to Sack state, -1 to remove.
void dfs_process_subtree(int u, int p, int val_sign) {
    update_sack_data(u, val_sign); // Process the current node 'u'
    for (int v : adj[u]) {        // For each child 'v' of 'u'
        if (v == p) continue;     // Skip parent
        dfs_process_subtree(v, u, val_sign); // Recursively process child's subtree
    }
}

// The main Sack (DSU on Tree) recursive function.
// This function efficiently calculates properties for all subtrees.
// 'u': current employee node (acting as the root of the subtree currently being aggregated).
// 'p': parent of 'u' in the DFS tree.
// 'keep': boolean flag. If true, the Sack state data for u's subtree should be preserved
//         after this function call (this is for the "big child" optimization).
//         If false, the data will be cleared.
void dfs_sack(int u, int p, bool keep) {
    // 1. Find the "big child" of 'u'.
    // The big child is the child with the largest subtree size.
    // Processing the big child last and keeping its data saves computations.
    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;
        }
    }

    // 2. Recursively call dfs_sack for all "small children" of 'u'.
    //    The 'keep' flag is set to 'false', so their data is cleared after their respective calls.
    //    The global Sack state arrays are effectively empty before processing the big child (or 'u' if no children).
    for (int v : adj[u]) {
        if (v == p || v == big_child) continue; // Skip parent and big_child
        dfs_sack(v, u, false); 
    }

    // 3. Recursively call dfs_sack for the "big child" (if one exists).
    //    The 'keep' flag is set to 'true'. This means after this call, the global Sack state arrays
    //    (freq_level, current_P_val, current_I_val) will hold the aggregated data
    //    for the big child's entire subtree.
    if (big_child != -1) {
        dfs_sack(big_child, u, true); 
    }

    // 4. Add employee 'u' (the current subtree root) itself to the Sack state.
    //    The Sack state now includes 'u' and, if a big child existed, its subtree data.
    update_sack_data(u, 1); 

    // 5. Add data from all "small children's" subtrees to the current Sack state.
    //    `dfs_process_subtree` is used here because we need to iterate through each small child's
    //    subtree and add all its nodes. The data for these small children was cleared earlier.
    for (int v : adj[u]) {
        if (v == p || v == big_child) continue; // Skip parent and big_child
        dfs_process_subtree(v, u, 1); // Add nodes of this small child's subtree
    }

    // 6. QUERY STEP:
    //    At this precise moment, the global Sack state arrays (freq_level, current_P_val, current_I_val)
    //    accurately reflect the aggregated information for the *entire subtree S_u* rooted at 'u'.
    //    We can now use these values to calculate the maximum team size (cand_P) and
    //    minimum switches (cand_S) if 'u' were chosen as the team lead.
    
    int k_u = emp_lang[u]; // The project's programming language is u's preferred language.
    
    // cand_P is the total number of people on the team if 'u' is lead and language is k_u.
    // This value is directly available from current_P_val[k_u], which has accumulated
    // sum_levels min(|S_u at l|, C[k_u][l]).
    long long cand_P = current_P_val[k_u]; 
    
    // cand_I is the number of initial members from S_u who prefer language k_u.
    // This is directly available from current_I_val[k_u].
    long long cand_I = current_I_val[k_u]; 
    
    // cand_S is the number of switches: total team members - initial team members.
    long long cand_S = cand_P - cand_I; 

    // Update the global best result (max_p_global, min_s_global)
    if (cand_P > max_p_global) { // If we found a larger team size
        max_p_global = cand_P;
        min_s_global = cand_S;
    } else if (cand_P == max_p_global) { // If team size is the same as current max
        min_s_global = min(min_s_global, cand_S); // Choose the one with fewer switches
    }
    
    // 7. CLEANUP STEP:
    //    If the 'keep' flag is 'false' (meaning 'u' is a small child of its caller,
    //    or 'u' is the root of the initial dfs_sack call from main),
    //    we must clear the data of u's entire subtree from the global Sack state arrays.
    //    This makes sure that when 'u's sibling (or parent's processing) continues,
    //    the Sack state is clean or reflects only the relevant parts.
    if (!keep) {
        // dfs_process_subtree with val_sign = -1 will remove all nodes in S_u
        // (which includes 'u' itself, its big child's subtree, and small children's subtrees
        // that were added in steps 4 and 5).
        dfs_process_subtree(u, p, -1); 
    }
}


// --- Main Function ---
int main() {
    // Speeds up C++ standard input/output operations.
    ios_base::sync_with_stdio(false); 
    cin.tie(NULL);                   

    // Read N (number of employees) and K (number of languages)
    cin >> n_nodes >> k_langs; 

    // Read preferred language for each employee (0 to N-1)
    for (int i = 0; i < n_nodes; ++i) {
        cin >> emp_lang[i]; 
    }
    
    // Build the directed tree from input. adj[boss] stores a vector of children.
    // Employee 0 is the CEO and has no boss.
    // Input describes bosses for employees 1 to N-1.
    vector<vector<int>> adj_build(n_nodes); // Temporary adjacency list for building
    for (int i = 1; i < n_nodes; ++i) { // For employee 'i' (from 1 to N-1)
        int b_val; // b_val is the boss of employee 'i'
        cin >> b_val; 
        adj_build[b_val].push_back(i); // Add 'i' as a child of its boss 'b_val'
    }
    // Copy the built tree structure to the global 'adj' array.
    for(int i=0; i<n_nodes; ++i) adj[i] = adj_build[i]; 


    // 1. Perform initial DFS traversals to gather basic tree information.
    dfs_levels(0, -1, 0); // Calculate levels for all nodes. CEO is node 0, has no parent (-1), at level 0.
    dfs_size(0, -1);      // Calculate subtree sizes for all nodes. Needed for Sack's big child heuristic.

    // 2. Precompute the sparse C[k][l] table (company-wide language counts per level).
    //    relevant_Cs_at_level[l] will store {lang, count} pairs for level l.
    relevant_Cs_at_level.resize(n_nodes); // Max level is n_nodes-1, so size n_nodes is sufficient.
    // Use a temporary map to aggregate counts before converting to vector of pairs.
    // temp_counts_for_Cs[level_idx] maps language_idx to its count at that level_idx.
    vector<map<int, int>> temp_counts_for_Cs(n_nodes); 
    for (int i = 0; i < n_nodes; ++i) { // For each employee
        if (level[i] < n_nodes) { // Level should always be < n_nodes
             temp_counts_for_Cs[level[i]][emp_lang[i]]++; // Increment count for their language at their level
        }
    }
    // Convert the map data to the final vector<vector<pair<int,int>>> structure
    // for efficient iteration in update_sack_data.
    for (int l = 0; l < n_nodes; ++l) { // For each level
        if (!temp_counts_for_Cs[l].empty()) { // If there are any employees at this level
            relevant_Cs_at_level[l].reserve(temp_counts_for_Cs[l].size()); // Pre-allocate memory (optimization)
            for (auto const& pair_kl : temp_counts_for_Cs[l]) { // For each {lang, count} at this level
                relevant_Cs_at_level[l].push_back({pair_kl.first, pair_kl.second});
            }
        }
    }
    
    // 3. Initialize Sack global state arrays.
    //    Global static arrays are typically zero-initialized by default in C++,
    //    but explicit initialization is good for clarity and safety.
    for(int i=0; i<k_langs; ++i) { // Max k_langs is 10^5, arrays are sized 100005
        current_P_val[i] = 0;
        current_I_val[i] = 0;
    }
    for(int i=0; i<n_nodes; ++i) { // Max n_nodes is 10^5, array is sized 100005
        freq_level[i] = 0;
    }
    
    // 4. Start the Sack algorithm from the CEO (employee 0).
    //    The 'keep' flag is 'false' for the initial call because we don't need to
    //    preserve the data of the entire tree after the process completes.
    dfs_sack(0, -1, false); 

    // 5. Output the final results.
    //    max_p_global will hold the maximum team size.
    //    min_s_global will hold the minimum switches for that maximum size.
    //    If N=0 was possible, max_p_global might be 0. But N>=1, so max_p_global>=1.
    //    min_s_global is initialized to INF and will be updated correctly.
    cout << max_p_global << " " << min_s_global << endl;

    return 0; // Indicate successful program execution
}

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