Submission #1145606

#	Time	Username	Problem	Language	Result	Execution time	Memory
1145606		monkey133	Team Coding (EGOI24_teamcoding)	C++20	23 / 100	4091 ms	44396 KiB

Main

#include <bits/stdc++.h>
using namespace std;
 
//––– Typedefs and macros –––
typedef int ll;
typedef pair<ll,ll> pi;
 
#define pb push_back
#define mp make_pair
 
// Use the same blk threshold (here 1000) as in your original code.
ll const blk = 1000;
 
//––– Global variables –––
const ll MAXN = 100005;
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, we record all depths encountered (later sorted & uniqued)
vector<ll> deplist[MAXN];
 
// Instead of unordered_map, we now use sorted vectors of (depth, frequency) pairs.
vector<pi> dp[MAXN], dp2[MAXN], dp3[MAXN], dp4[MAXN], dp5[MAXN];
// For each color (indexed 0..k-1) we store the overall frequency by depth.
vector< vector<pi> > colMapVec; 
 
// Our final answer is stored as (res, -swaps) so that max() works lexicographically.
pi ans = mp(-1000000000, -1000000000);
 
//–––– Helper routines for sorted frequency arrays –––––
 
// getValue: binary search for key in sorted vector v; returns frequency if found, else 0.
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;
}
 
// addValue: 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));
}
 
// mergeVectors: merge two sorted frequency vectors (dest and src) summing counts for equal keys.
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 < dest.size() && j < 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 < dest.size()){
        merged.pb(dest[i]);
        i++;
    }
    while(j < src.size()){
        merged.pb(src[j]);
        j++;
    }
    dest = move(merged);
}
 
//–––– DFS #1: Euler Tour and depth recording –––––
void precomp1(ll x = 0) {
    in[x] = ++label;
    for (ll u : adj[x]) {
        dep[u] = dep[x] + 1;
        precomp1(u);
    }
    out[x] = label;
    // Record the depth for this node’s color.
    deplist[color[x]].pb(dep[x]);
}
 
//–––– DFS #2: DSU on Tree for building frequency arrays –––––
void precomp2(ll x = 0) {
    // Update the global frequency table for color[color[x]].
    addValue(colMapVec[color[x]], dep[x], 1);
    for (ll u : adj[x]) {
        precomp2(u);
        // Merge dp[u] into dp[x] using small-to-large merging.
        if (dp[x].size() < dp[u].size()) swap(dp[x], dp[u]);
        mergeVectors(dp[x], dp[u]);
        dp[u].clear();
    }
    addValue(dp[x], dep[x], 1);
    // For “small” colors, also build dp2[x] (only at the depths that appear for that color).
    if (!big[color[x]]) {
        for (ll d : deplist[color[x]]) {
            int freq = getValue(dp[x], d);
            if (freq > 0)
                addValue(dp2[x], d, freq);
        }
    }
}
 
//–––– DFS for “big” colors –––––
// In this DFS we merge dp4 and dp5 arrays for later query.
void dfs(ll x, ll c, ll yes = 1) {
    for (ll u : adj[x]) {
        dfs(u, c, yes & (color[x] != c));
        if (dp4[x].size() < dp4[u].size()) swap(dp4[x], dp4[u]);
        mergeVectors(dp4[x], dp4[u]);
        if (dp5[x].size() < dp5[u].size()) swap(dp5[x], dp5[u]);
        mergeVectors(dp5[x], dp5[u]);
        dp4[u].clear();
        dp5[u].clear();
    }
    if (yes && color[x] == c) {
        int res = 0, res2 = 0;
        // Iterate over unique depths for color c that are deeper than dep[x].
        for (ll d : deplist[c]) {
            if (d <= dep[x]) continue;
            int a = getValue(dp4[x], d);
            int b = getValue(colMapVec[c], d);
            int take = min(a, b);
            res += take;
            int cval = getValue(dp5[x], d);
            res2 += max(0, take - cval);
        }
        ans = max(ans, mp(res, -res2));
    }
    addValue(dp4[x], dep[x], 1);
    if (c == color[x])
        addValue(dp5[x], dep[x], 1);
}
 
//–––– Process “small” colors –––––
void solve(ll c) {
    int sz = nodelist[c].size();
    // For every pair (x,y) of nodes of color c with x an ancestor of y, update dp3.
    for (int i = 0; i < sz; i++) {
        for (int j = i + 1; j < sz; j++) {
            ll 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);
        }
    }
    // Then, for each node x of color c, compute candidate answers using dp2 and dp3.
    for (int i = 0; i < sz; i++) {
        ll x = nodelist[c][i];
        int res = 0, res2 = 0;
        for (ll d : deplist[c]) {
            if (d <= dep[x]) continue;
            int a = getValue(dp2[x], d);
            int b = getValue(colMapVec[c], d);
            int take = min(a, b);
            res += take;
            int cval = getValue(dp3[x], d);
            res2 += max(0, take - cval);
        }
        dp2[x].clear();
        dp3[x].clear();
        ans = max(ans, mp(res, -res2));
    }
}
 
//–––– Main –––––
int main(){
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);
    
    cin >> n >> k;
    colMapVec.resize(k);  // prepare the global frequency table for each color
 
    for (ll i = 0; i < n; i++){
        cin >> color[i];
        cnt[color[i]]++;
        nodelist[color[i]].pb(i);
    }
    for (ll i = 0; i < k; i++){
        if (cnt[i] > blk)
            big[i] = 1;
    }
    for (ll i = 1; i < n; i++){
        cin >> par[i];
        adj[par[i]].pb(i);
    }
    
    precomp1();
    // For each color, sort and unique the list of depths.
    for (ll 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());
    }
    precomp2();
    
    // Process each color:
    // For “big” colors we use a DFS that uses dp4 and dp5;
    // for “small” colors we use the solve() routine.
    for (ll i = 0; i < k; i++){
        if (big[i]) {
            // Clear temporary arrays for the DFS.
            for (ll j = 0; j < n; j++){
                dp5[j].clear();
                dp4[j].clear();
            }
            dfs(0, i, 1);
        } else {
            solve(i);
        }
    }
    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