Submission #1145612

#	Time	Username	Problem	Language	Result	Execution time	Memory
1145612		monkey133	Team Coding (EGOI24_teamcoding)	C++20	62 / 100	4126 ms	782396 KiB

Main

#include <bits/stdc++.h>
using namespace std;
 
//–––– Custom hash (unchanged) ––––
struct custom_hash {
    static uint64_t splitmix64(uint64_t x) {
        x += 0x9e3779b97f4a7c15ULL;
        x = (x ^ (x >> 30)) * 0xbf58476d1ce4e5b9ULL;
        x = (x ^ (x >> 27)) * 0x94d049bb133111ebULL;
        return x ^ (x >> 31);
    }
    size_t operator()(uint64_t x) const {
        static const uint64_t random_constant = chrono::steady_clock::now().time_since_epoch().count();
        return splitmix64(x + random_constant);
    }
};
 
//–––– Typedefs and macros ––––
typedef int ll;
typedef pair<ll,ll> pi;
 
#define pb push_back
#define mp make_pair
 
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];
vector<ll> deplist[MAXN];  // for each color, record all depths
 
// We use vectors of unordered_maps (with custom_hash) for our DP arrays.
vector<unordered_map<ll,ll, custom_hash>> dp(MAXN), dp2(MAXN), dp3(MAXN), dp4(MAXN), dp5(MAXN), colMap(MAXN);
 
// Answer stored as (res, -swaps) so that max() works lexicographically.
pi ans = mp(-1000000000, -1000000000);
 
//–––– precomp1: Euler Tour and record depths ––––
void precomp1(ll x = 0) {
    in[x] = ++label;
    for (ll u : adj[x]) {
        dep[u] = dep[x] + 1;
        precomp1(u);
    }
    out[x] = label;
    // Instead of inserting into a set, just push_back.
    deplist[color[x]].pb(dep[x]);
}
 
//–––– precomp2: DSU–on–Tree merging using unordered_map ––––
void precomp2(ll x = 0) {
    colMap[color[x]][dep[x]]++;  // update global frequency for this color
    for (ll u : adj[x]) {
        precomp2(u);
        // Merge dp[u] into dp[x]: swap if dp[x] is smaller
        if(dp[x].size() < dp[u].size())
            swap(dp[x], dp[u]);
        // Pre–reserve: add size of dp[u] to dp[x] to avoid rehashing overhead.
        dp[x].reserve(dp[x].size() + dp[u].size());
        for(auto it = dp[u].begin(); it != dp[u].end(); ++it)
            dp[x][it->first] += it->second;
        dp[u].clear();
    }
    dp[x][dep[x]]++;
    // For “small” colors, also build dp2[x].
    if (!big[color[x]]) {
        for (ll d : deplist[color[x]])
            dp2[x][d] += dp[x][d];
    }
}
 
//–––– dfs: process “big” colors using DSU–merging on dp4 and dp5 ––––
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]);
        dp4[x].reserve(dp4[x].size() + dp4[u].size());
        for(auto it = dp4[u].begin(); it != dp4[u].end(); ++it)
            dp4[x][it->first] += it->second;
 
        if(dp5[x].size() < dp5[u].size())
            swap(dp5[x], dp5[u]);
        dp5[x].reserve(dp5[x].size() + dp5[u].size());
        for(auto it = dp5[u].begin(); it != dp5[u].end(); ++it)
            dp5[x][it->first] += it->second;
 
        dp4[u].clear();
        dp5[u].clear();
    }
    if(yes && color[x] == c) {
        ll res = 0, res2 = 0;
        // Loop only over unique depths for color c.
        for (ll d : deplist[c]) {
            if(d <= dep[x])
                continue;
            ll a = dp4[x][d];
            ll b = colMap[c][d];
            ll take = a < b ? a : b;
            res += take;
            ll cval = dp5[x][d];
            res2 += take > cval ? take - cval : 0;
        }
        ans = max(ans, mp(res, -res2));
    }
    dp4[x][dep[x]]++;
    if(c == color[x])
        dp5[x][dep[x]]++;
}
 
//–––– solve: process “small” colors ––––
void solve(ll c) {
    int sz = nodelist[c].size();
    // For every pair (x,y) with x an ancestor of y among nodes of color c.
    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])
                dp3[x][dep[y]]++;
        }
    }
    // For each node x of color c, combine dp2[x] and dp3[x] to update answer.
    for (int i = 0; i < sz; i++) {
        ll x = nodelist[c][i];
        ll res = 0, res2 = 0;
        for (ll d : deplist[c]) {
            if(d <= dep[x])
                continue;
            ll a = dp2[x][d];
            ll b = colMap[c][d];
            ll take = a < b ? a : b;
            res += take;
            ll cval = dp3[x][d];
            res2 += take > cval ? take - cval : 0;
        }
        dp2[x].clear();
        dp3[x].clear();
        ans = max(ans, mp(res, -res2));
    }
}
 
//–––– Main ––––
int main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);
    
    cin >> n >> k;
    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: use DFS for “big” colors; use solve() for “small” colors.
    for (ll i = 0; i < k; i++){
        if (big[i]) {
            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: 12 / 12
• Subtask 2: 0 / 19
• Subtask 3: 27 / 27
• Subtask 4: 23 / 23
• Subtask 5: 0 / 19