PROBLEM LINKS
DIFFICULTY
MEDIUM
EXPLANATION
Make the tree a rooted tree, choosing any vertex arbitrarily as the root. We use dynamic programming to solve this problem.
We store 2 things for every vertex :
- ways1 ( VERTEX, pathLength ) -> Number of subtrees ( using vertices from subtree rooted at VERTEX ) with the longest path ( with VERTEX as one endpoint ) less than or equal to pathLength and diameter less than k.
- ways2 ( VERTEX, pathLength ) -> Number of subtrees ( using vertices from subtree rooted at VERTEX ) with the longest path ( with VERTEX as one endpoint ) equal to pathLength and diameter less than k.
The recurrence relation:
- ways1( VERTEX, pathLength ) = ways2 ( VERTEX,pathLength ) + ways1( VERTEX, pathLength-1 )
- To calculate ways2( VERTEX, pathLength ), notice that atleast one of the child node should have pathLength-1 as the length of the longest path with the child node as endpoint and other child node should have the pathLength as pathLength-1 OR k-pathLength-1( diameter cannot exceed k), whichever is minimum. So, we take a loop of i on the first child node with length of Longest path as pathLength-1.
ways2(VERTEX,pathLength) = Summation of [ Product { j < i } [ ways1(jth child node,MIN{pathLength-2,k-1-pathLength} ] * Product{j > i}[ways1(jth child node, MIN(pathLength-1,k-1-pathLength)] * ways2(ith child node,pathLength-1) ] over all child node as i.
Notice that Product{j < i} component is the prefix product and Product {j > i} is the suffix product. Both values can be accessed in O(1) after O(n) preprocessing.
It is easy to see that the actual answer is the summation of ways1(VERTEX, k)-1 [ -1 to exclude the empty subtree] for all vertices.