Tags: math, combinatorics, DSU
Now, this is a really beautiful problem from AtCoder: you are given an array of $N$ balls and two numbers $X,Y$. $col[i], w[i]$ represent the color and weight of $i^{th}$ ball respectively. You can swap two balls,$i$ and $j$ of same color if $w[i]+w[j] \le X$; you can swap two balls,$i$ and $j$ of different color if $w[i]+w[j] \le Y$. You need to find the total number of distinct permutations of colors that you can obtain.
The solution is as follows:
The key concept is that if you can swap balls $A,B$ and balls $B,C$, then you can also swap $A,C$ regardless of the condition. Why?
1. Let the original position be $A ... B ... C$
2. Swap $A,B$ to obtain $B ... A ... C$
3. Swap $B,C$ to obtain $C ... A ... B$
4. Swap $A,B$ to obtain $C ... B ... A$
Now, this is a really beautiful problem from AtCoder: you are given an array of $N$ balls and two numbers $X,Y$. $col[i], w[i]$ represent the color and weight of $i^{th}$ ball respectively. You can swap two balls,$i$ and $j$ of same color if $w[i]+w[j] \le X$; you can swap two balls,$i$ and $j$ of different color if $w[i]+w[j] \le Y$. You need to find the total number of distinct permutations of colors that you can obtain.
The solution is as follows:
The key concept is that if you can swap balls $A,B$ and balls $B,C$, then you can also swap $A,C$ regardless of the condition. Why?
1. Let the original position be $A ... B ... C$
2. Swap $A,B$ to obtain $B ... A ... C$
3. Swap $B,C$ to obtain $C ... A ... B$
4. Swap $A,B$ to obtain $C ... B ... A$
Voila! You can simply swap $A,C$.
This means that we need to make clusters of balls. A cluster contains balls that can be swapped into each other's position and the total answer would be the product of total permutations for each cluster of balls.
So the only problem is to make clusters. Each ball can be connected to $n-1$ other balls. So we have around $O(n^2)$ edges that we need to see and add. We need something faster.
Note that each edge can be of two types, connecting balls of: same color or different color. We use DSU to add edges between the balls.
Note that each edge can be of two types, connecting balls of: same color or different color. We use DSU to add edges between the balls.
1. If for every color $C$, we connect the lightest ball of color $C$ to all the balls of same color which can be joined as per the rules, we completely incorporate all the edges of type that connect two balls of same color.
2. Now we find the two balls of different colors and minimum weight: let them be A,B.
3. Connect all the balls of color different from $col[A]$ to $A$ if we can connect them.
4. Now connect balls with color same as $col[A]$ to $B$ if we can connect them.
Each step takes $O(n)$ time.
How this works?
Consider $i^{th}, j^{th}$ balls, and assume that they can be swapped. We will prove that our idea of addng edges will already collect both these balls in the same cluster.
Case 1: Same Color
Let the lightest ball of that color be $k$.
Since they can be swapped,
$w[i]+w[j]\le X$
This means that $w[i]+w[k] \le X, w[j]+w[k] \le X$
$\therefore $ our idea of connecting ball to the lightest ball of that color makes the cluster complete and we never miss out any edge of this type.
Case 2: Different Color
You can work(similar to what we did above) out to see that if $w[i]+w[j] \le Y$, then our idea would also collect both these balls in the same cluster.
How this works?
Consider $i^{th}, j^{th}$ balls, and assume that they can be swapped. We will prove that our idea of addng edges will already collect both these balls in the same cluster.
Case 1: Same Color
Let the lightest ball of that color be $k$.
Since they can be swapped,
$w[i]+w[j]\le X$
This means that $w[i]+w[k] \le X, w[j]+w[k] \le X$
$\therefore $ our idea of connecting ball to the lightest ball of that color makes the cluster complete and we never miss out any edge of this type.
Case 2: Different Color
You can work(similar to what we did above) out to see that if $w[i]+w[j] \le Y$, then our idea would also collect both these balls in the same cluster.
Now we have made clusters, all that is left is to compute answers for each and then multiply them.
Check the code for implementation of above idea.
Overall, I had too much fun solving this problem and I learnt something from it.
Hey rachit!
ReplyDeleteFor the different color case isn't it sufficient to just take the global minimum of all minimums and make the edges.Why take two?
Please correct if i am missing something
What about the balls of same color as the global minimum ball? They also need their edges to connect to a ball of different color, right?
DeleteBut initially in step 1 we have considered edges for same colours. So,if the minimum of that colour(say color X) doesn't have an edge with its same colour balls in the DSU different colours will also not have the edge to that color(X).
DeleteI don't know what makes you belief that. Or maybe you are missing the different constraints: swapping balls of same color must weight < X, while swapping balls of different color must weight < Y.
DeleteHope you can see it now why its necessary.
believe*
Delete