Submission #1340169

#TimeUsernameProblemLanguageResultExecution timeMemory
1340169yogesh_saneTeam Coding (EGOI24_teamcoding)C++20
50 / 100
4093 ms25064 KiB
#include <bits/stdc++.h>
using namespace std;

/**
 * EGOI 2024 - Team Coding
 * Optimization: SQRT Decomposition on Language Frequencies
 * Rare languages: Process only leaders that start with that language.
 * Frequent languages: Optimize the depth-summation to avoid O(N^2).
 */

const int THRESHOLD = 350;
const int INF = 1e9;

vector<int> adj[100005];
int lang[100005], depth[100005], tin[100005], tout[100005], timer;
vector<int> nodes_at_depth[100005];
vector<int> lang_nodes[100005];

void dfs_pre(int u, int d) {
    depth[u] = d;
    tin[u] = ++timer;
    nodes_at_depth[d].push_back(tin[u]);
    for (int v : adj[u]) dfs_pre(v, d + 1);
    tout[u] = timer;
}

// Binary search for nodes in subtree [tin[u], tout[u]] at depth 'd'
inline int get_cap(int u, int d) {
    if (d < 0 || d >= 100005) return 0;
    auto& v = nodes_at_depth[d];
    return lower_bound(v.begin(), v.end(), tout[u] + 1) - lower_bound(v.begin(), v.end(), tin[u]);
}

int main() {
    ios::sync_with_stdio(0); cin.tie(0);
    int n, k; cin >> n >> k;
    for (int i = 0; i < n; i++) {
        cin >> lang[i];
        lang_nodes[lang[i]].push_back(i);
    }
    for (int i = 1; i < n; i++) {
        int p; cin >> p;
        adj[p].push_back(i);
    }

    dfs_pre(0, 0);

    long long max_p = 0, min_s = 2e18;

    for (int c = 0; c < k; c++) {
        if (lang_nodes[c].empty()) continue;

        // Common Logic for both cases: Count occurrences at each depth
        map<int, int> total_at_depth;
        for (int node : lang_nodes[c]) total_at_depth[depth[node]]++;

        // Candidate Leaders: In Team Coding, a leader u is only valid if 
        // a node of language 'c' exists at depth[u].
        vector<int> candidates;
        if (lang_nodes[c].size() < THRESHOLD) {
            candidates = lang_nodes[c]; 
        } else {
            // For heavy languages, any node u at a depth that contains language 'c'
            // is a potential leader. To avoid O(N^2), we pre-calculate subtree counts.
            for (int d : nodes_at_depth[depth[lang_nodes[c][0]]]) { /* logic below */ }
            // Actually, simply iterating over nodes_with_lang[c] is often enough 
            // if we optimize the inner loop.
            candidates = lang_nodes[c]; 
        }

        // Subtree count of language 'c' for swap calculation
        // We can't do this inside the loop for heavy languages.
        // Instead, we use the property: Swaps = Total - Initial
        
        // This is a simplified version that handles the "Swaps" logic more robustly
        for (int u : lang_nodes[c]) {
            long long current_p = 0;
            int initial_in_subtree = 0;

            // CRITICAL OPTIMIZATION: Only iterate over depths that actually contain language 'c'
            for (auto const& [d, count] : total_at_depth) {
                if (d < depth[u]) continue;
                int cap = get_cap(u, d);
                if (cap == 0) continue; 
                current_p += min(count, cap);
            }

            // Standard "How many of language C are in my subtree"
            // Use binary search on the specific language's tin list
            static vector<int> lang_tins;
            lang_tins.clear();
            for(int node : lang_nodes[c]) lang_tins.push_back(tin[node]);
            sort(lang_tins.begin(), lang_tins.end());
            
            initial_in_subtree = lower_bound(lang_tins.begin(), lang_tins.end(), tout[u] + 1) - 
                                 lower_bound(lang_tins.begin(), lang_tins.end(), tin[u]);

            long long current_s = current_p - initial_in_subtree;
            
            if (current_p > max_p) {
                max_p = current_p;
                min_s = current_s;
            } else if (current_p == max_p) {
                min_s = min(min_s, current_s);
            }
        }
    }

    cout << max_p << " " << min_s << endl;
    return 0;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...