Submission #1340170

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

const int MAXN = 100005;
const int THRESHOLD = 320;

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

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;
}

// Fast Subtree Capacity Query
inline int get_cap(int u, int d) {
    if (d < 0 || d >= n) 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);
    if (!(cin >> n >> k)) return 0;

    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;
    long long min_s = 2e18;

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

        // 1. Pre-calculate global supply per depth for this language
        vector<pair<int, int>> depth_supply;
        map<int, int> temp_map;
        for (int node : lang_nodes[c]) temp_map[depth[node]]++;
        for (auto const& [d, cnt] : temp_map) depth_supply.push_back({d, cnt});

        // 2. Pre-calculate how many of language 'c' are in each node's subtree
        // We do this once per language to avoid O(N^2)
        static int subtree_lang_cnt[MAXN];
        // Only reset nodes that were affected or use a timestamp to stay O(N)
        auto compute_subtree = [&](auto self, int u) -> int {
            int cnt = (lang[u] == c ? 1 : 0);
            for (int v : adj[u]) cnt += self(self, v);
            return subtree_lang_cnt[u] = cnt;
        };

        // For HEAVY languages, we pre-calculate subtree counts using one DFS
        if (lang_nodes[c].size() >= THRESHOLD) {
            compute_subtree(compute_subtree, 0);
        }

        for (int u : lang_nodes[c]) {
            long long current_p = 0;
            
            // Optimization: Iterate over relevant depths only
            for (auto const& p_depth : depth_supply) {
                int d = p_depth.first;
                int supply = p_depth.second;
                if (d < depth[u]) continue;
                
                int cap = get_cap(u, d);
                if (cap == 0) continue; 
                current_p += min(supply, cap);
            }

            int initial_in_subtree;
            if (lang_nodes[c].size() >= THRESHOLD) {
                initial_in_subtree = subtree_lang_cnt[u];
            } else {
                // For LIGHT languages, just count manually
                initial_in_subtree = 0;
                for (int node : lang_nodes[c]) {
                    if (tin[node] >= tin[u] && tin[node] <= tout[u]) initial_in_subtree++;
                }
            }

            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...