Submission #1340746

#TimeUsernameProblemLanguageResultExecution timeMemory
1340746toast12Team Coding (EGOI24_teamcoding)C++20
100 / 100
3499 ms30360 KiB
#include <bits/stdc++.h>
using namespace std;

const int SQRTN = 316;

int n;
vector<vector<int>> adj, pos;
vector<int> color, depth, en, ex;
vector<vector<int>> depth_en, depth_ex;

int cur_time = 0;

pair<int, int> subtask3(int color);
pair<int, int> subtask2(int color);

void dfs(int cur, int p) {
    depth[cur] = depth[p]+1;
    en[cur] = cur_time;
    depth_en[depth[cur]].push_back(en[cur]);
    for (auto e : adj[cur]) {
        cur_time++;
        dfs(e, cur);
    }
    ex[cur] = cur_time;
    depth_ex[depth[cur]].push_back(ex[cur]);
}

vector<int> roots;

int main() {
    int k;
    cin >> n >> k;
    adj.resize(n);
    pos.resize(k);
    color.resize(n);
    for (int i = 0; i < n; i++) {
        int x;
        cin >> x;
        pos[x].push_back(i);
        color[i] = x;
    }
    for (int i = 1; i < n; i++) {
        int p;
        cin >> p;
        adj[p].push_back(i);
    }
    depth.resize(n);
    en.resize(n);
    ex.resize(n);
    depth_en.resize(n);
    depth_ex.resize(n);
    depth[0] = -1;
    dfs(0, 0);
    pair<int, int> ans = {0, INT_MAX};
    for (int i = 0; i < k; i++) {
        pair<int, int> res;
        // few people with one color --> subtask 3
        if (pos[i].size() <= SQRTN) res = subtask3(i);
        else { // many people with one color --> subtask 2
            res = subtask2(i);
            roots.clear();
        }
        if (res.first > ans.first) {
            ans.first = res.first;
            ans.second = res.second;
        }
        else if (res.first == ans.first) ans.second = min(ans.second, res.second);
    }
    cout << ans.first << ' ' << ans.second << '\n';
    return 0;
}

bool subtree(int u, int v) {
    if (en[u] >= en[v] && ex[u] <= ex[v]) return true;
    return false;
}

pair<int, int> subtask3(int color) {
    vector<int> nodes = pos[color];
    map<int, int> depth_freq;
    for (auto node : nodes) depth_freq[depth[node]]++;
    pair<int, int> res = {0, INT_MAX};
    for (auto root : nodes) {
        // try every node with this color as the root
        int cnt = 0, swaps = 0;
        map<int, int> subtree_freq_depth;
        for (auto node : nodes) {
            if (subtree(node, root)) {
                cnt++;
                subtree_freq_depth[depth[node]]++;
            }
        }
        for (auto node : nodes) {
            if (!subtree(node, root)) {
                // this node can only be included if it can been swapped into the root's subtree
                // find the total number of nodes at this node's depth in the root's subtree
                auto start = lower_bound(depth_en[depth[node]].begin(), depth_en[depth[node]].end(), en[root])-depth_en[depth[node]].begin();
                auto end = upper_bound(depth_ex[depth[node]].begin(), depth_ex[depth[node]].end(), ex[root])-depth_ex[depth[node]].begin();
                int tot = end-start;
                // find the number of "unoccupied" nodes at this node's depth in the root's subtree
                if (tot > 0 && tot - subtree_freq_depth[depth[node]] > 0) { // there is at least one node that can be swapped
                    cnt++, swaps++;
                    subtree_freq_depth[depth[node]]++;
                }
            }
        }
        if (cnt > res.first) {
            res.first = cnt;
            res.second = swaps;
        }
        else if (cnt == res.first) res.second = min(res.second, swaps);
    }
    return res;
}

void dfs2(int cur, int p, int c) {
    if (color[cur] == c) {
        roots.push_back(cur);
        return;
    }
    for (auto e : adj[cur]) dfs2(e, cur, c);
}

vector<pair<int, int>> f;

void dfs3(int cur, int p, int c) {
    if (color[cur] == c) f[depth[cur]].first++;
    f[depth[cur]].second++;
    for (auto e : adj[cur]) dfs3(e, cur, c);
}

pair<int, int> subtask2(int color) {
    vector<int> freq(n);
    for (auto x : pos[color]) freq[depth[x]]++;
    // only the topmost nodes of this color will be viable roots
    dfs2(0, 0, color);
    pair<int, int> res = {0, INT_MAX};
    for (auto root : roots) {
        int cnt = 1, swaps = 0;
        f.clear();
        f.resize(n);
        dfs3(root, root, color);
        for (int j = depth[root]+1; j < n; j++) {
            if (f[j].second == 0) break;
            cnt += f[j].first;
            int outside = freq[j]-f[j].first;
            cnt += min(f[j].second-f[j].first, outside);
            swaps += min(f[j].second-f[j].first, outside);
        }
        if (cnt > res.first) {
            res.first = cnt;
            res.second = swaps;
        }
        else if (cnt == res.first) res.second = min(res.second, swaps);
    }
    return res;
}
#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...