PROBLEM LINK :
Author : Asim Krishna Prasad
Tester : Rupesh Maity
DIFFICULTY :
MEDIUM-HARD
PREREQUISITES :
Trees, Probability, Expected Values.
PROBLEM :
Given a tree and number of marbles at each node of the tree. You are given the fact that each marble rolls from it’s own node to some other node. The node at which the marble stops can be any node except for it’s own node , with equal probability. You are to find the expected number of edges from which all the marbles will pass.
EXPLANATION :
The first thing to note is, the way first sample is solved in the problem statement is just to distract the users. One needs to understand the linearity of expected value to solve this problem.
Using linearity of expected values we can come to the conclusion that, the answer is : summation of { probability that a edge will be beautiful } * 1.
We multiply by one because each edge is of length 1.
One can get to the same conclusion by observing the way expected value is calculated.
expected value = summation of { probability of a set of beautiful edges * number of beautiful edges in this set}
For each of the edges in the input, it will contribute a probability for each of the sets that contain it, which is, the probability of this edge being beautiful.
Now the problem reduces to, calculate the probability that an edge will be traveled by all the marbles.
Let us consider an edge A to B. Now we can disconnect the tree by removing this edge from the tree.
We have now got two subtrees A and B.
For edge A - B to be traveled by all the marbles, the condition is, every marble in the subtree A should move to subtree B and every marble in subtree B should move to subtree A.
Let us say number of marbles in subtree A is ma and number of marbles in subtree B is mb.
Similarly , number of nodes in subtree A is na and number of nodes in subtree B is nb.
Let us define foroneA as the probability for any marble in subtree A to select any node in subtree B.
So, foroneA = nb / (nodes - 1)
(nodes - 1) because it could have selected any node except for it’s own.
Similarly, foroneB = na / (nodes - 1)
and hence the probability of edge A - B being traveled by all the marbles, or edge A - B being beautiful is : (foroneAma) * (foroneBmb).
Summation of this for all the edges gives us the required answer.
Let us start with an O(n^2) solution.
- Iterate through each edge.
- For every edge A-B call dfs at node A such that it doesn't go at subtree of B.
- The dfs should return
- The number of nodes in subtree A.
- The total number of marbles in subtree A.
- From this we can calculate number of nodes in subtree B and number of marbles in subtree B using the fact that number of marbles remains constant throughout the process.
- We use the same formula stated above ( (foroneAma) * (foroneBmb) ) and add it to the answer.
To reduce the complexity to O(n).
- Root the tree at any node.
- Precompute the sizes of subtrees and number of marbles in that subtree.
- Use the same method as above but you don't need the dfs for every edge as you have the values precomputed.