Submission #1206637

#	Time	Username	Problem	Language	Result	Execution time	Memory
1206637		ocean	Team Coding (EGOI24_teamcoding)	C++17	23 / 100	4096 ms	75352 KiB

Main

#include <iostream>
#include <vector>
#include <unordered_map> // Use unordered_map for potentially faster average time
#include <algorithm>     // For std::min, std::max, std::swap

using namespace std;

// Global variables for tree structure and properties
vector<vector<int>> adj; // Adjacency list for children (representing the hierarchy)
vector<int> languages;   // Preferred language for each employee (indexed by employee ID)
vector<int> depth_arr;   // Depth of each employee from Anneke (CEO 0 at depth 0)

// Global variables to store the best results found
int max_P = 0; // Maximum number of employees in the project team
int min_S = 0; // Minimum number of switches needed to achieve max_P

// Global variable to store total counts of (language, depth) pairs for the entire company
// total_depth_lang_counts[d][k] stores count of employees at depth 'd' with language 'k'
vector<unordered_map<int, int>> total_depth_lang_counts;


// DFS to compute the depth of each employee in the tree
// u: current employee ID
// d: current depth
void compute_depths_dfs(int u, int d) {
    depth_arr[u] = d; // Set depth for current employee
    for (int v : adj[u]) { // Iterate through subordinates (children)
        compute_depths_dfs(v, d + 1); // Recurse for children at next depth
    }
}

// DFS function that uses small-to-large merging to compute subtree statistics
// It returns a pointer to a map representing (depth -> (language -> count)) for the subtree rooted at 'u'.
// The caller takes ownership of the returned pointer and is responsible for its deletion or merging.
unordered_map<int, unordered_map<int, int>>* dfs_solve(int u) {
    // Create a new map for this node's direct contribution
    unordered_map<int, unordered_map<int, int>>* u_subtree_counts = new unordered_map<int, unordered_map<int, int>>();
    // Add the current node's language at its depth to its subtree counts
    (*u_subtree_counts)[depth_arr[u]][languages[u]]++;

    // Process children
    for (int v : adj[u]) {
        // Recursively get subtree counts for child 'v'
        unordered_map<int, unordered_map<int, int>>* v_subtree_counts = dfs_solve(v);

        // Small-to-large merging logic: always merge the smaller map into the larger one.
        // This optimizes the overall time complexity.
        if (u_subtree_counts->size() < v_subtree_counts->size()) {
            swap(u_subtree_counts, v_subtree_counts); // 'u_subtree_counts' now points to the larger map
        }

        // Merge 'v_subtree_counts' into 'u_subtree_counts'
        for (auto const& [d_val, lang_map_at_d] : *v_subtree_counts) { // Iterate through depths in smaller map
            for (auto const& [lang_val, count] : lang_map_at_d) {     // Iterate through languages at each depth
                (*u_subtree_counts)[d_val][lang_val] += count;       // Add count to the larger map
            }
        }
        delete v_subtree_counts; // Free memory of the merged (smaller) map
    }

    // Now, 'u_subtree_counts' contains the complete (depth -> (language -> count))
    // statistics for the subtree rooted at 'u'.
    // Calculate P and S if 'u' is chosen as the team lead.
    int current_team_lead_lang = languages[u]; // Project language is team lead's preferred language
    int p_for_u = 0;                          // Initial team members within u's subtree with target language
    int s_for_u = 0;                          // Switches needed for u as team lead

    // Iterate through all depths present in u's subtree
    for (auto const& [d, lang_map_at_d_const] : *u_subtree_counts) {
        int current_lang_count_in_subtree = 0;
        if (lang_map_at_d_const.count(current_team_lead_lang)) {
            current_lang_count_in_subtree = lang_map_at_d_const.at(current_team_lead_lang);
        }
        p_for_u += current_lang_count_in_subtree;

        // Calculate 'X_d': employees in u's subtree at depth 'd' NOT preferring target_lang
        int in_subtree_wrong_lang_at_d = 0;
        for (auto const& [lang, count] : lang_map_at_d_const) {
            if (lang != current_team_lead_lang) {
                in_subtree_wrong_lang_at_d += count; // These can be swapped out
            }
        }
        
        // Calculate 'Y_d': employees NOT in u's subtree at depth 'd' who DO prefer target_lang
        int total_same_lang_at_d_globally = 0;
        // Check if the language exists at this depth globally before accessing
        if (total_depth_lang_counts[d].count(current_team_lead_lang)) {
            total_same_lang_at_d_globally = total_depth_lang_counts[d].at(current_team_lead_lang);
        }
        
        // Employees in u's subtree with target language
        int in_subtree_same_lang_at_d = current_lang_count_in_subtree; 
        int out_of_subtree_correct_lang_at_d = total_same_lang_at_d_globally - in_subtree_same_lang_at_d;

        // Add to switches: min(X_d, Y_d) possible swaps at this depth
        s_for_u += min(in_subtree_wrong_lang_at_d, out_of_subtree_correct_lang_at_d);
    }

    int current_P = p_for_u + s_for_u; // Total team size for u as team lead
    int current_S = s_for_u;           // Total switches for u as team lead

    // Update global maximum P and minimum S
    if (current_P > max_P) {
        max_P = current_P;
        min_S = current_S;
    } else if (current_P == max_P) { // If same P, choose the one with fewer switches
        min_S = min(min_S, current_S);
    }

    return u_subtree_counts; // Return the calculated subtree counts for 'u'
}


int main() {
    // Optimize C++ standard streams for faster input/output
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int N, K; // N: number of employees, K: number of programming languages
    cin >> N >> K; //[cite: 21]

    languages.resize(N);
    for (int i = 0; i < N; ++i) {
        cin >> languages[i]; // Read preferred language for each employee [cite: 22]
    }

    // Build adjacency list. Employee 'i' reports to 'b_i', so 'b_i' is parent of 'i'. [cite: 23]
    adj.resize(N);
    for (int i = 1; i < N; ++i) { // Employees 1 to N-1 have bosses [cite: 24]
        int b;
        cin >> b; // Direct boss of employee 'i' [cite: 23]
        adj[b].push_back(i); // Add 'i' as a child of 'b'
    }

    // Compute depths for all employees, starting with Anneke (CEO 0) at depth 0 [cite: 5]
    depth_arr.resize(N);
    compute_depths_dfs(0, 0);

    // Precompute total counts of (language, depth) for the entire company.
    // Max possible depth is N-1.
    total_depth_lang_counts.resize(N); 
    for (int i = 0; i < N; ++i) {
        total_depth_lang_counts[depth_arr[i]][languages[i]]++;
    }

    // Start the main DFS to find the maximum P and minimum S.
    // The pointer returned for the root (employee 0) needs to be deleted to clean up memory.
    unordered_map<int, unordered_map<int, int>>* final_subtree_counts_root = dfs_solve(0);
    delete final_subtree_counts_root; // Clean up memory

    // Output the results [cite: 25]
    cout << max_P << " " << min_S << endl;

    return 0;
}

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