Submission #1146130

#	Time	Username	Problem	Language	Result	Execution time	Memory
1146130		monkey133	Team Coding (EGOI24_teamcoding)	C++20	23 / 100	4094 ms	36320 KiB

Main

#include <bits/stdc++.h>
using namespace std;
 
//–––––– Typedefs and Macros ––––––
typedef int ll;  // adjust if needed (using int for indices)
typedef pair<ll, ll> pi;
 
#define pb push_back
#define mp make_pair
 
const ll MAXN = 100005;
 
//–––––– Global Variables ––––––
ll n, k, label = 0;
ll color[MAXN], par[MAXN], in[MAXN], out[MAXN], cnt[MAXN], dep[MAXN], big[MAXN];
vector<ll> adj[MAXN], nodelist[MAXN];
// For each color, record all encountered depths (to later sort & unique)
vector<ll> deplist[MAXN];
 
// Instead of unordered_maps, we now store frequency tables as sorted vectors of (depth, frequency) pairs.
vector<vector<pi>> dp(MAXN), dp2(MAXN), dp3(MAXN), dp4(MAXN), dp5(MAXN);
// Global frequency table for each color (by depth). Index: color
vector<vector<pi>> colMapVec;
 
// Final answer stored as (res, -swaps) so that lexicographical max works.
pi ans = mp(-1000000000, -1000000000);
 
//–––––– Helper Functions for Sorted Frequency Tables –––––//
 
// Binary search: return frequency for key in sorted vector v, or 0 if not found.
inline int getValue(const vector<pi>& v, int key) {
    int lo = 0, hi = (int)v.size() - 1;
    while (lo <= hi) {
        int mid = (lo + hi) >> 1;
        if (v[mid].first == key)
            return v[mid].second;
        else if (v[mid].first < key)
            lo = mid + 1;
        else
            hi = mid - 1;
    }
    return 0;
}
 
// Insert or update key in sorted vector v by adding val.
inline void addValue(vector<pi>& v, int key, int val) {
    auto it = std::lower_bound(v.begin(), v.end(), mp(key, 0));
    if (it != v.end() && it->first == key)
        it->second += val;
    else
        v.insert(it, mp(key, val));
}
 
// Merge two sorted frequency vectors: merge src into dest.
inline void mergeVectors(vector<pi>& dest, const vector<pi>& src) {
    vector<pi> merged;
    merged.reserve(dest.size() + src.size());
    int i = 0, j = 0;
    while (i < (int)dest.size() && j < (int)src.size()){
        if (dest[i].first == src[j].first) {
            merged.pb(mp(dest[i].first, dest[i].second + src[j].second));
            i++; j++;
        } else if (dest[i].first < src[j].first) {
            merged.pb(dest[i]);
            i++;
        } else {
            merged.pb(src[j]);
            j++;
        }
    }
    while (i < (int)dest.size()){
        merged.pb(dest[i]); i++;
    }
    while (j < (int)src.size()){
        merged.pb(src[j]); j++;
    }
    dest = move(merged);
}
 
//–––––– Iterative DFS for Euler Tour (Computes in[], out[], dep[]) –––––//
void iterativeDFS_precomp1() {
    vector<int> st;
    vector<int> ptr(n, 0); // pointer for each node's child index
    dep[0] = 0;
    st.pb(0);
    while (!st.empty()) {
        int u = st.back();
        // First time visiting u: record entry time and add its depth to deplist.
        if (ptr[u] == 0) {
            in[u] = ++label;
            deplist[color[u]].pb(dep[u]);
        }
        if (ptr[u] < (int)adj[u].size()) {
            int v = adj[u][ptr[u]++];
            dep[v] = dep[u] + 1;
            st.pb(v);
        } else {
            out[u] = label;
            st.pop_back();
        }
    }
}
 
//–––––– Get Postorder Order (Iterative) –––––//
vector<int> getPostorder() {
    vector<int> post;
    vector<int> st;
    st.pb(0);
    while(!st.empty()){
        int u = st.back();
        st.pop_back();
        post.pb(u);
        for (int v : adj[u])
            st.pb(v);
    }
    reverse(post.begin(), post.end());
    return post;
}
 
//–––––– Precomp2: DSU-on-Tree Merging Using Sorted Vectors (Iterative Postorder) –––––//
void precomp2_iterative() {
    // Initialize global frequency table for each color.
    colMapVec.assign(k, vector<pi>());
    for (int u = 0; u < n; u++) {
        addValue(colMapVec[color[u]], dep[u], 1);
    }
    vector<int> post = getPostorder();
    // Process nodes in postorder (children before parent).
    for (int u : post) {
        // Merge each child's dp into dp[u].
        for (int v : adj[u]) {
            if (dp[u].size() < dp[v].size())
                swap(dp[u], dp[v]);
            mergeVectors(dp[u], dp[v]);
            dp[v].clear();
        }
        // Add the current node.
        addValue(dp[u], dep[u], 1);
        // For “small” colors, build dp2[u] using only the depths that appear for that color.
        if (!big[color[u]]) {
            for (int d : deplist[color[u]]) {
                int freq = getValue(dp[u], d);
                if (freq > 0)
                    addValue(dp2[u], d, freq);
            }
        }
    }
}
 
//–––––– Solve for “big” colors using Iterative Postorder Merging for dp4/dp5 –––––//
void solve_big_iterative(int c) {
    // Build a "yes" array: yes[u] is true if none of u's ancestors (excluding u) has color c.
    vector<bool> yes(n, false);
    {
        vector<int> st;
        vector<int> ptr(n, 0);
        yes[0] = true;
        st.pb(0);
        while(!st.empty()){
            int u = st.back();
            if(ptr[u] < (int)adj[u].size()){
                int v = adj[u][ptr[u]++];
                yes[v] = yes[u] && (color[u] != c);
                st.pb(v);
            } else {
                st.pop_back();
            }
        }
    }
    vector<int> post = getPostorder();
    // Clear dp4 and dp5 for all nodes.
    for (int u = 0; u < n; u++){
        dp4[u].clear();
        dp5[u].clear();
    }
    // Process nodes in postorder.
    for (int u : post) {
        if (yes[u] && color[u] == c) {
            int res = 0, res2 = 0;
            // For every unique depth for color c deeper than dep[u], compute contributions.
            for (int d : deplist[c]) {
                if(d <= dep[u]) continue;
                int a = getValue(dp4[u], d);
                int b = getValue(colMapVec[c], d);
                int take = (a < b ? a : b);
                res += take;
                int cval = getValue(dp5[u], d);
                res2 += (take > cval ? take - cval : 0);
            }
            ans = max(ans, mp(res, -res2));
        }
        // Add the current node u.
        addValue(dp4[u], dep[u], 1);
        if(c == color[u])
            addValue(dp5[u], dep[u], 1);
        // Merge u's dp4/dp5 into its parent (if not root).
        if(u != 0) {
            int p = par[u];
            if (dp4[p].size() < dp4[u].size())
                swap(dp4[p], dp4[u]);
            mergeVectors(dp4[p], dp4[u]);
            dp4[u].clear();
            if (dp5[p].size() < dp5[u].size())
                swap(dp5[p], dp5[u]);
            mergeVectors(dp5[p], dp5[u]);
            dp5[u].clear();
        }
    }
}
 
//–––––– Solve for “small” colors –––––//
void solve_small(int c) {
    int sz = (int)nodelist[c].size();
    // For every pair of nodes of color c (where one is an ancestor of the other),
    // update dp3 for the ancestor.
    for (int i = 0; i < sz; i++) {
        for (int j = i + 1; j < sz; j++) {
            int x = nodelist[c][i], y = nodelist[c][j];
            if (in[x] > in[y])
                swap(x, y);
            if (in[x] <= in[y] && in[y] <= out[x])
                addValue(dp3[x], dep[y], 1);
        }
    }
    // For each node of color c, combine dp2 and dp3 to update the candidate answer.
    for (int i = 0; i < sz; i++) {
        int x = nodelist[c][i];
        int res = 0, res2 = 0;
        for (int d : deplist[c]) {
            if(d <= dep[x]) continue;
            int a = getValue(dp2[x], d);
            int b = getValue(colMapVec[c], d);
            int take = (a < b ? a : b);
            res += take;
            int cval = getValue(dp3[x], d);
            res2 += (take > cval ? take - cval : 0);
        }
        dp2[x].clear();
        dp3[x].clear();
        ans = max(ans, mp(res, -res2));
    }
}
 
//–––––– Main Function –––––//
int main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);
    
    cin >> n >> k;
    // Resize global frequency table vector for each color.
    colMapVec.resize(k);
    
    // Read node colors and fill nodelist and frequency counts.
    for (int i = 0; i < n; i++){
        cin >> color[i];
        cnt[color[i]]++;
        nodelist[color[i]].pb(i);
    }
    // Mark "big" colors: here using a threshold (blkThreshold).
    const int blkThreshold = 316;  // can adjust threshold as needed
    for (int i = 0; i < k; i++){
        if (cnt[i] > blkThreshold)
            big[i] = 1;
    }
    // Read parent pointers and build tree. (Assuming input is 0-indexed.)
    for (int i = 1; i < n; i++){
        cin >> par[i];
        adj[par[i]].pb(i);
    }
    
    // Compute Euler Tour iteratively: fill in[], out[], dep[] and deplist.
    iterativeDFS_precomp1();
    // For each color, sort and unique the deplist.
    for (int i = 0; i < k; i++){
        sort(deplist[i].begin(), deplist[i].end());
        deplist[i].erase(unique(deplist[i].begin(), deplist[i].end()), deplist[i].end());
    }
    // Build dp and dp2 using DSU-on-tree merging (iterative postorder).
    precomp2_iterative();
    
    // Process each color: if color is "big" use iterative DFS merging (dp4/dp5), else use solve_small.
    for (int c = 0; c < k; c++){
        if (big[c]) {
            solve_big_iterative(c);
        } else {
            solve_small(c);
        }
    }
    // Output final answer (note: answer is stored as (res, -swaps) so we output ans.first+1 and -ans.second).
    cout << ans.first + 1 << " " << -ans.second << "\n";
    return 0;
}

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