Distinct Colors - CSES

I was trying to solve a question on cses - Distinct Colors, My solution based on the book Competitive Programmer’s Handbook by Antti Laaksonen page number 170(Merging data structures)

#include<bits/stdc++.h>
using namespace std;
const int Mxn = 2e5;
unordered_set<int> adj[Mxn],mp[Mxn];
int c[Mxn],ans[Mxn],n;
 
void dfs(int u=0, int p = -1){
    for(auto v: adj[u]){
        if(v != p){
            dfs(v,u);
            if(mp[v].size() > mp[u].size())
                swap(mp[v],mp[u]);
            mp[u].insert(mp[v].begin(),mp[v].end());
        }
    }
    mp[u].insert(c[u]);
    ans[u] = mp[u].size();
}
 
int main(){
    cin>>n;
    for(int i=0;i<n;i++) cin>>c[i];
    for(int i=1;i<n;i++){
        int x,y; cin>>x>>y; x--,y--;
        adj[x].insert(y);
        adj[y].insert(x);
    }
    dfs();
    for(int i=0;i<n;i++) cout<<ans[i]<<' ';
}

What I want to know is that when I swap the two maps in the above code using

if(mp[v].size() > mp[u].size())
    swap(mp[v],mp[u]);

The code gets Accepted, but when I do not swap these maps I get a TLE, here is a quote from the book :

The merging of nodes can be done as follows: We go through the children of s and at each child u merge d[s] and d[u]. We always copy the contents from d[u] to d[s]. However, before this, we swap the contents of d[s] and d[u] if d[s] is smaller than d[u]. By doing this, each value is copied only O(log(n)) times during the tree traversal, which ensures that the algorithm is efficient.

how does swapping the maps according to their size ensures that each element would be at most copied O(log(N)) times?

If all the colors are unique, then the complexity to copy one map to another would-be O(N) right? That would make total complexity O(N^2).

1 Like

I am blown away by the explanation. Thank you so much !! :bowing_man: