제출 #1146126

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
1146126		monkey133	Team Coding (EGOI24_teamcoding)	C++20	23 / 100	4094 ms	36356 KiB

Main

#include <bits/stdc++.h>
using namespace std;
 
//–––––– Typedefs and macros ––––––
typedef int ll; // using int for simplicity (adjust if needed)
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 (will sort & unique later)
vector<ll> deplist[MAXN];
 
// We now use vectors of (key,frequency) pairs instead of unordered_map.
vector<vector<pi>> dp(MAXN), dp2(MAXN), dp3(MAXN), dp4(MAXN), dp5(MAXN);
// Global frequency table per color (indexed by color) – we'll fill this in precomp2.
vector<vector<pi>> colMapVec; // size will be k
 
// Global answer stored as (res, -swaps) so that lexicographical max works.
pi ans = mp(-1000000000, -1000000000);
 
//–––––– Helper functions for sorted frequency arrays –––––//
 
// 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 (precomp1) –––––//
void iterativeDFS_precomp1() {
    vector<int> stack;
    vector<int> ptr(n, 0);
    // Start at root (node 0); set depth 0.
    dep[0] = 0;
    stack.push_back(0);
    while(!stack.empty()){
        int u = stack.back();
        // First time visiting u: record in[u] 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;
            stack.push_back(v);
        } else {
            out[u] = label;
            stack.pop_back();
        }
    }
}
 
//–––––– Compute Postorder Order Iteratively –––––//
vector<int> getPostorder() {
    vector<int> post;
    vector<int> stack;
    vector<bool> visited(n, false);
    stack.push_back(0);
    while(!stack.empty()){
        int u = stack.back();
        stack.pop_back();
        post.pb(u);
        for (int v : adj[u])
            stack.push_back(v);
    }
    reverse(post.begin(), post.end());
    return post;
}
 
//–––––– Precomp2: DSU-on-Tree merging using sorted vectors (iterative, postorder based) –––––//
void precomp2_iterative() {
    // Initialize colMapVec for each color.
    colMapVec.assign(k, vector<pi>());
    // For each node, update global frequency table for its color.
    for (int u = 0; u < n; u++) {
        addValue(colMapVec[color[u]], dep[u], 1);
    }
    // Get postorder order.
    vector<int> post = getPostorder();
    // Process nodes in postorder (children before parent).
    for (int u : post) {
        // For each child of u, merge 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 self.
        addValue(dp[u], dep[u], 1);
        // For small colors, build dp2[u] = frequencies for depths in deplist[color[u]]
        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 DFS postorder for DP4 and DP5 –––––//
void solve_big_iterative(int c) {
    // Compute yes[u] for each node: yes[u] = true if none of u's ancestors (excluding u) have color equal to c.
    vector<bool> yes(n, false);
    {
        vector<int> stack;
        vector<int> ptr(n, 0);
        yes[0] = true;  // root always true.
        stack.push_back(0);
        while(!stack.empty()){
            int u = stack.back();
            if(ptr[u] < (int)adj[u].size()){
                int v = adj[u][ptr[u]++];
                yes[v] = yes[u] && (color[u] != c);
                stack.push_back(v);
            } else {
                stack.pop_back();
            }
        }
    }
    // Get postorder order.
    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 the path from root to u is "clean" (yes[u] true) and u's color equals c, update answer.
        if(yes[u] && color[u] == c) {
            int res = 0, res2 = 0;
            // Iterate over all unique depths for color c.
            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);
            }
            if(mp(res, -res2) > ans)
                ans = mp(res, -res2);
        }
        // Add self to dp4[u] and dp5[u].
        addValue(dp4[u], dep[u], 1);
        if(c == color[u])
            addValue(dp5[u], dep[u], 1);
        // Merge current node u 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 (same as before) –––––//
void solve_small(int c) {
    int sz = (int)nodelist[c].size();
    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 (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();
        if(mp(res, -res2) > ans)
            ans = mp(res, -res2);
    }
}
 
//–––––– Main –––––//
int main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);
    
    cin >> n >> k;
    // Resize colMapVec to hold k colors.
    colMapVec.resize(k);
    
    // Read node colors and build nodelist and count frequencies.
    for (int i = 0; i < n; i++){
        cin >> color[i];
        cnt[color[i]]++;
        nodelist[color[i]].pb(i);
    }
    // Mark "big" colors: if count > blk threshold.
    const int blkThreshold = 316;  // or 1000 as in your original code; here we use 316
    for (int i = 0; i < k; i++){
        if (cnt[i] > blkThreshold)
            big[i] = 1;
    }
    // Read parent pointers and build tree.
    // Root is assumed to be 0.
    for (int i = 1; i < n; i++){
        cin >> par[i];
        // Assuming parent input is 0-indexed; if 1-indexed, do: par[i]--;
        adj[par[i]].pb(i);
    }
    
    // Iterative DFS to compute Euler Tour (in, out, dep) and fill 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());
    }
    // DSU merging precomputation: build dp and dp2 and global colMapVec.
    precomp2_iterative();
    
    // Process each color: for big colors, use iterative DFS (dp4/dp5) merging;
    // for small colors, use the solve_small routine.
    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)
    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