PROBLEM LINK:
DIFFICULTY:
Medium
PREREQUISITES:
DFS, trees, probablities
PROBLEM:
Given a weighted tree of N vertices, some vertices are special, you will be given a list of length M containing the labels of special vertices. if we pick randomly an ordered pair of special vertices what is the mathematical expectation of the distance of the vertices of the chosen pair? output the answer as irreducible fraction.
EXPLANATION:
There are M2 different ordered pairs of special vertices, and since all pairs equiprobable to be chosen then we just need to find the sum of distances of all ordered pairs of special vertices - let’s call it S- then divide it by M2 and it will be the required answer, since the problem asks to output the answer as irreducible fraction then we should find the gcd of S and M2 then divide both numbers by it then output them.
Calculating the sum of distances of all ordered pairs of special vertices:
a straight forward idea is to do a DFS/BFS from each of the special vertices to calculate its distance from each of the other special vertices, unfortunately, this is not fast enough to get full score for this problem so we should search for a better idea.
note that S is sum of M2 paths lengths, each path is a sum of lengths of edges on that path. so if we break each path into the edges that the path consist of, then S will be the sum of edges lengths, each edge length is repeated in the sum number of times equal to the number of paths that passes through it.
let’s take the second sample input as an example:
6 6
1 3 1
2 3 2
3 4 3
4 5 4
4 6 5
1 2 3 4 5 6
- the first edge has length 1 and there are 10 paths which pass through it (note that back and forth of same route is counted twice)
- the second edge has length 2 and there are 10 paths which pass through it
- the third edge has length 3 and there are 18 paths which pass through it
- the fourth edge has length 4 and there are 10 paths which pass through it
- the fifth edge has length 5 and there are 10 paths which pass through it
that means S is equal to 1 * 10 + 2 * 10 + 3 * 18 + 4 * 10 + 5 * 10 = 174 and if we divide it by 36 (M2) it will be same answer as the output.
so now all we need to find a way to count for each edge the number of paths which start and end in a special vertex and pass through that edge, if we could count that then we are done.
Counting the number of paths that pass through each node
Let’s observe an edge of some tree, if we delete it the tree will be separated into exactly two trees(possibly a tree of a single vertex) each tree was connected from one side of that edge, now what are the paths which were passing through the deleted edge? it’s clear that they are the paths which start at a special vertex from one of the two trees formed after deleting that edge and ends at a special vertex from the other tree.
so using some basic combinatorics rules, the number of paths which pass through the edge is 2 * A * B, where A is the number of special vertices in the first tree and B is the number of special vertices in the second tree, the number 2 comes from the fact that edge path is counted twice (back and forth)
now all we need is find a way to count for each edge the number of special vertices on each of the two sides of tree. let’s choose an arbitrary vertex (say vertex 1) to be the root of the tree, now every vertex has a parent (except the root) and has its own sub-tree. counting the number of special vertices in a sub-tree is the same of counting number of special vertices of one side of the tree of the edge which connect the root of that sub-tree with its parent, the other side of the tree is simply M - number of special vertices in the first side.
counting the number of special vertices of every sub-tree can be done with a single DFS, here’s pseudo code of the DFS
ans = 0;
dfs(int v){
if (v is special vertex){
sum[v] = 1
} else {
sum[v] = 0
}
for each child u of v{
dfs(u)
sum[v] = sum[v] + sum[u];
}
ans = ans + 2 * sum[v] * (M - sum[v])
}
in the code sum[v] means the number of special vertices in the sub-tree of vertex v, as you can see the DFS works by first initiating sum[v] to be 1 if v itself is special vertex otherwise it will be 0, after that we recursively call the dfs function for each child of v in order to compute the number of special vertices of its sub-tree, after that we add this number to sum[v].
in the end, we apply the formula 2 * A * B mentioned above by adding to the final answer the quantity 2 * sum[v] * (M - sum[v]).
SETTER’S SOLUTION
Can be found here.
TESTER’S SOLUTION
Can be found here.