Submission #1206651

#	Time	Username	Problem	Language	Result	Execution time	Memory
1206651		ocean	Team Coding (EGOI24_teamcoding)	C++17	54 / 100	4091 ms	40304 KiB

Main

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

using namespace std;

const long long INF = 4e18; // Using a large enough value for infinity

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] (total employees of lang k at level l)
// relevant_Cs_at_level[l] stores a list of pairs {lang_k, count}, 
// where count = C[lang_k][l] and count > 0.
vector<vector<pair<int, int>>> relevant_Cs_at_level;

// Global arrays for Sack states
// current_P_val[k_test] = sum_l min(freq_level[l], C[k_test][l]) for current Sack bag
// current_I_val[lang_x] = count of employees with lang_x in current Sack bag
long long current_P_val[100005]; 
long long current_I_val[100005]; 
// freq_level[l] = count of employees at level l in current Sack bag
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];
    }
}

// Adds/removes node_x effects from global Sack state variables
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 all k_test based on the change of freq_level[lvl_x]
    // Only languages k_test for which C[k_test][lvl_x] > 0 will have their min term potentially change.
    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;
            // Contribution from min(old_freq, C_val)
            if (old_freq_at_lvl_x > 0) { 
                term_change -= min(old_freq_at_lvl_x, C_val_k_test_lvl_x);
            }
            // Contribution from min(new_freq, C_val)
            if (new_freq_at_lvl_x > 0) { 
                term_change += min(new_freq_at_lvl_x, C_val_k_test_lvl_x);
            }
            // Apply the net change to current_P_val for this k_test
            current_P_val[k_test] += term_change;
        }
    }
    
    // Update current_I_val for the language of the node being processed
    current_I_val[lang_x] += val_sign;
}

// Helper to add/remove all nodes in u's subtree from Sack state
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);
    }
}

// Main Sack algorithm DFS
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;
        }
    }

    // Process small children first, and clear their data from Sack state
    for (int v : adj[u]) {
        if (v == p || v == big_child) continue;
        dfs_sack(v, u, false); 
    }

    // Process big child (if any), keeping its data in Sack state
    if (big_child != -1) {
        dfs_sack(big_child, u, true); 
        // Sack state now reflects big_child's subtree
    }

    // Add node u itself to the Sack state
    // (If there was a big child, its state is already there; if not, state is empty)
    update_sack_data(u, 1); 

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

    // Now Sack state (freq_level, current_P_val, current_I_val) reflects subtree S_u
    // Calculate P_u and S_u for u as team lead
    int k_u = emp_lang[u]; // Project language is language of team lead u
    long long cand_P = current_P_val[k_u];
    long long cand_S = cand_P - current_I_val[k_u]; // Switches = Total Team - Initial Team

    // Update global best
    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 data for u's subtree is not to be kept (i.e., u is a small child or root of full Sack call)
    if (!keep) {
        // Remove all nodes of S_u (which includes u and all its descendants) from Sack state
        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); // Compute levels for all nodes

    // Precompute sparse C[k][l] into relevant_Cs_at_level
    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) { // level[i] is already computed and < n_nodes
             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 state arrays (static global arrays are zero-initialized by default in C++)
    // Explicitly zeroing can be done for safety/clarity if needed, but not strictly necessary for global/static.
    /*
    for(int i=0; i < k_langs && i < 100005; ++i) { 
        current_P_val[i] = 0;
        current_I_val[i] = 0;
    }
    for(int i=0; i < n_nodes && i < 100005; ++i) { 
        freq_level[i] = 0;
    }
    */
    
    dfs_size(0, -1); // Compute subtree sizes needed for Sack's big child heuristic
    
    dfs_sack(0, -1, false); // Start Sack algorithm from root (node 0)

    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