Submission #1145409

#	Time	Username	Problem	Language	Result	Execution time	Memory
1145409		monkey133	Team Coding (EGOI24_teamcoding)	C++20	0 / 100	1546 ms	1114112 KiB

Main

#include <bits/stdc++.h>
using namespace std;

typedef int ll;
const ll blk = 316;
const int 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];

// Instead of an unordered_set for deplist (indexed by color), we use a vector
vector<ll> deplist[MAXN];

// We change the dp arrays from fixed global arrays of unordered_map to vectors of unordered_map.
// (This way, if n < MAXN, we are not “wasting” space.)
vector<unordered_map<ll,ll>> dp(MAXN), dp2(MAXN), dp3(MAXN), dp4(MAXN), dp5(MAXN), colVec(MAXN);

// The answer is stored as a pair; note that we update the second element as negative so that
// using max() gives the lexicographically maximum pair.
pair<ll,ll> ans = make_pair(-1000000000, -1000000000);

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 an unordered_set, just push_back the depth.
    deplist[color[x]].push_back(dep[x]);
}

void precomp2(ll x = 0) {
    // For every node x, record (in a map per color) the frequency of each depth.
    colVec[color[x]][dep[x]]++;
    for (ll u : adj[x]) {
        precomp2(u);
        // Merge dp[u] into dp[x] using the “small-to–large” trick.
        if (dp[x].size() < dp[u].size()) swap(dp[x], dp[u]);
        for (auto &p : dp[u])
            dp[x][p.first] += p.second;
        dp[u].clear(); // free memory from child's dp after merging
    }
    dp[x][dep[x]]++;
    // Build dp2 for node x – only for depths that appear for the color of x.
    for (ll d : deplist[color[x]])
        dp2[x][d] += dp[x][d];
}

void dfs(ll x, ll c, ll yes = 1) {
    for (ll u : adj[x]) {
        dfs(u, c, yes & (color[x] != c));
        // Merge dp4[u] into dp4[x]
        if (dp4[x].size() < dp4[u].size()) swap(dp4[x], dp4[u]);
        for (auto &p : dp4[u])
            dp4[x][p.first] += p.second;
        dp4[u].clear();
        // Merge dp5[u] into dp5[x]
        if (dp5[x].size() < dp5[u].size()) swap(dp5[x], dp5[u]);
        for (auto &p : dp5[u])
            dp5[x][p.first] += p.second;
        dp5[u].clear();
    }
    if (yes && color[x] == c) {
        ll res = 0, res2 = 0;
        // Iterate only over the unique depths for color c.
        for (ll d : deplist[c]) {
            if (d <= dep[x]) continue;
            ll take = min(dp4[x][d], colVec[c][d]);
            res += take;
            res2 += max(0, take - dp5[x][d]);
        }
        ans = max(ans, make_pair(res, -res2));
    }
    dp4[x][dep[x]]++;
    if (c == color[x])
        dp5[x][dep[x]]++;
}

void solve(ll c) {
    int sz = nodelist[c].size();
    // Precompute dp3 for nodes of color c.
    for (int i = 0; i < sz; i++) {
        ll x = nodelist[c][i];
        for (int j = i + 1; j < sz; j++) {
            ll 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 (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 take = min(dp2[x][d], colVec[c][d]);
            res += take;
            res2 += max(0, take - dp3[x][d]);
        }
        ans = max(ans, make_pair(res, -res2));
    }
    // Clear dp3 for nodes of color c to free memory.
    for (int i = 0; i < sz; i++) {
        ll x = nodelist[c][i];
        dp3[x].clear();
    }
}

int main(){
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    cin >> n >> k;
    for (ll i = 0; i < n; i++){
        cin >> color[i];
        cnt[color[i]]++;
        nodelist[color[i]].push_back(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]].push_back(i);
    }
    
    precomp1();
    // For each color, sort and unique–ify the deplist vector so that we don’t iterate duplicate 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();
    
    for (ll i = 0; i < k; i++){
        if (big[i]) {
            // For every DFS run we clear the dp4/dp5 arrays for all nodes
            for (ll j = 0; j < n; j++){
                dp4[j].clear();
                dp5[j].clear();
            }
            dfs(0, i, 1);
        } else {
            solve(i);
        }
        // Free memory no longer needed for this color.
        colVec[i].clear();
        deplist[i].clear();
    }
    
    cout << ans.first + 1 << " " << -ans.second << "\n";
    return 0;
}

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