C. Experience

A CodeForces EDU Problem based on Union-Find

This blog is about the problem Experience, which is the third practice problem on DSU in the ITMO Academy pilot course available from Codeforces EDU.

Here's a link to the problem (although you would need to be logged in to CF for this to work). The statement is reproduced below for your reference.

📖

Problem Statement

In a novel online game, the players fight the monsters and get the experience, as usual. To fight monsters players join together in raid teams. After the destruction of the monster, all the players of the team get the same amount of experience points. The special feature of the game is that a team cannot be split up and no one can leave a team. The only supported operation is to join two teams together. Since there are already a lot of people playing the game, you are asked to maintain the experience points of the players.

Input

The first line of the input contains two integers n and m  — the number of players and the number of queries.

Next m lines contain the description of queries, one per line. A query can be of three types:

  • join  X Y — join two teams to which players X and Y belong to (if they are already in the same team, nothing changes).
  • add X V — add  (1 ≤ V ≤ 100) experience points to each player in a team to which player  belongs to,
  • get X — output the current number of experience points of player X.

Initially, each player has 0 experience points and each of the player is in its own team of size one.

Output

For each query get X output the current number of experience points of player X on a separate line

The problem seems fairly simple at first glance, we perform union operation for the query of type 'join', and for every 'add' query we iterate through all the elements to which 'X' belongs to and update their score. And finally for the query 'get X', fetch and display the score of X.

It is also intuitive to not perform path compression since a player is only entitled to the team scores that were added after he joined that team.

Let's look at the (pseudo!) code for the 'add' operation.

int find_set (int i) {                       
	        if (i == parent [i]) // No path compression 
							return i; 
        return find_set (parent[i]); 
    }
// we initially have a vector<int> score(n+1) 
// assigned with all values 0;

void updateScore (int X, int V) {
    int index = find_set (family);
		for (int i=0; i<parent.size(); i++) {
				if ( find_set (parent [i]) == index){
							score [i] += V;
				}
	  }
}

While there is nothing incorrect with this approach, the constraints of the problem are too tight to pass solution in the given time-limit. Observe that in the worst case, the 'updateScore' method will have to iterate over all the sets for (~10^5 times). And as established, this solution actually gives the verdict 'Time Limit Exceeded'.

At this point, it seems a little obvious that we have to somehow update the score only for the parent set and keep track of the extra score a player is assigned after it merges with a set.

For example, consider the sample test case:

3 6
add 1 100
join 1 3
add 1 50
get 1
get 2
get 3

For simplicity, we will consider three players 1,2,3.

The first query, 'add 1 100' means that a singleton set of Player 1 has a score of 100 now.

In the second query, Player 1 and 3 merge to form a team, now Player 3 is a part of a set with score 100 but it is not entitled to the score which Player 1 received in the first query since it wasn't a team member at that point so it's only fair that we assign the Player 3 with a score of 0 and keep track of its extra-score by incrementing it to 100.

In the third query we add a score of 50 to Player 3 and hence, Player 1 since they are from the same team. A manual calculation makes it obvious that the score of Player 1 is now 150 and Player 3 is 50 and Player 2 is obviously at score 0.

The important thing to note here is that since we assigned an extra score of 100 to Player 3, we can easily find the final score of Player 3 by subtracting the extra-score of Player 3 from the final score of the set to which the player belongs, this will have to be done recursively if there are multiple members in the set.

A simple code, with commented explanation wherever necessary, using a DSU implementation as close as possible to the DSU defined in the lectures is as follows:

#include <bits/stdc++.h>
using namespace std;
#define vi  vector<int>


class UnionFind {
    private:
        vi p, rank, setSize, score, extraScore;
    public:

        // Initialization

        UnionFind(int N){
            // The goal here is to create N singleton sets.

            p.resize(N);
            // To begin with, everyone is their own parent:
            for(int i = 0; i < N; i++)
                p[i] = i;

            // the sizes of the individual sets are one:
            setSize.assign(N, 1);

            // sizes of score and assigning 0 to all players
            score.assign(N,0);

            // sizes of extra score and assigning 0 to all players
            extraScore.assign(N,0);
            
        }

        // FindSet(i): return the representative element of the set that i belongs to.

        int findSet(int i){
            // if already at a representative element,
	          // signified by the fact that the parent pointer 
						// points to the element itself, then return.

            if(p[i] == i)
                return i;

            // otherwise find the representative of the parent, 
						// without path compression
            else
                return findSet(p[i]);
        }

        // isSameSet(i,j): returns true if and only if i and j 
				// belong to the same set.

        bool isSameSet(int i, int j){
            // Observe that i and j belong to the same set
            // if and only if they have the same representative.
            return findSet(i) == findSet(j); 
        }

        
    
        void unionSet(int i, int j){
            if(isSameSet(i,j)) 
                return; 
            
            int x = findSet(i);
            int y = findSet(j);

            if(setSize[x] > setSize[y])
                swap(x,y);
            
            p[x] = y; 

            
            setSize[y] += setSize[x];

            // store the extra score that the smaller set is not entitled to

            extraScore[x] = score[y];

        }


        void updateScore(int x, int v){
          
          // finding the parent of the player 
					// and updating its score

          x = findSet(x);
          score[x] += v;                             
       }

        int findScore(int x){

        // if X is its own parent, (this is the base 
				// case implying that at this point the player 
				// X joined/created this team) 

         if(x==p[x]) 
            return score[x];
         
        // this is a bit tricky as we have to recursively 
        // find the base case, in every step we subtract 
        // the extra score the player X has from when he 
				// was not a part of this team
         else 
            return (score[x] - extraScore[x]) + findScore(p[x]);                      
         }
};

int x,y,v;
int main(){
    std::ios_base::sync_with_stdio(0);
		std::cin.tie(0);std::cout.tie(0);

    int n, m;
    cin >> n >> m;
    
    UnionFind UF(n+1); 

    while(m--){
        string S;
        cin >> S;      

        if(S == "join"){
            cin >> x >> y;            
            UF.unionSet(x,y);
        }
        else if (S == "add") {
            cin >> x >> v;
            UF.updateScore(x,v);
        }
        else
        {   
             cin >> x;
             cout << UF.findScore(x) << '\n';
         }
    }

    return 0;
}

The constraints of the problem are actually a little too tight, so I would suggest removing non essential variables like rank or number of sets to get an AC in case any one is having issues passing certain test cases.

The above implementation in its current form passes all the test cases easily and gets an 'ACCEPTED' verdict.

📝

This post was contributed by @Faraz. Join the Discord community for the course to discuss this!