PROBLEM LINKS
DIFFICULTY
EASY
Solution: Tree traversal Complexity: O(N)
EXPLANATION
Consider the rooted tree T, with the 1 as the root. Now the required node C for a pair (A,B) is the lowest common ancestor of A and B in this tree. Direct implementation of a text book algorithm leads to a O(N^2logN) or a O(N^2) solution to the problem depending on whether its O(logN) lca or O(1) lca algorithm. Since N <= 1e4, and there are 25 test cases, we need something better than O(N^2) to pass.
Consider a non-leaf node u in the tree. Now we have to compute the number of pairs of nodes whose lca is u. Let u have k children c1,c2,ā¦,ck. Now the subtree rooted at each child c _ i is independent. For all combinations of vertices (a,b) with a in c _ i and b in c _ j , i ! = j , have u as their lca. So we must find 2 * (|c1| * |c2| + |c1| * |c3| + ā¦ |c1| * |ck| + |c2| * |c3| + |c2| * c4 + ā¦ + |ck| * |ck-1|) , ie. all C(k,2) combinations of |c _ i|s where |c _ i| represents the number of nodes in the subtree rooted at c _ i. The two factor is because each pair (a,b) appears in the matrix twice as (a,b) and (b,a).
Now to compute this, there are multiple approaches:
1.The sum is equal to (sum of |c _ i|)^2 minus sum of |c _ i|^2
An alternate and neater approach to compute this sum is to keep the partial sum of c _ iās until a point and to multiply the next c value to the present partial sum. ie.
long long partialsum = 0 , res = 0;
for ( int i=1;i<=k;i++ )
{
res += partialsum * C[i];
partialsum += C[i];
}
res * 2 is the required sum. Now all the pairs of the form (u,some node in subtree rooted at u) are not considered as we are only considering the pairs among the children of u. So we must add twice the size of subtree rooted at u (excluding u) to the solution. (u,u) also appears in the matrix, so this is another entry for Im(u). So Im(u) = res * 2 + 2 * (size of subtree rooted at u) + 1.
Note: If partialsum is initialized to 1, then the subtree term can be ignored completely, since the size of subtree rooted at u excluding u is equal to sum of |c _ i|s.
Once the Im values are calculated, the required sum S can easily be computed.