PROBLEM LINK:
Author: Sergey Kulik
Tester: Roman Furko
Editorialist: Balajiganapathi Senthilnathan
DIFFICULTY:
Medium
PREREQUISITES:
Dynamic programming
PROBLEM:
A heavy light decomposition of a tree is where you label exactly one edge going out of each node (i.e. away from the root) as heavy and others as light. A correct heavy light decomposition is where the number of light edges from root to any node is at most \lfloor{log_2} n \rfloor. Given a tree, find the number of correct heavy light decomposition it can have.
QUICK EXPLANATION:
This can solved by using dynamic programming with the state(node index, number of light edges seen from root so far). To calculate dp[x][cnt] we loop through each outgoing edge of x and mark it as heavy in turn. We recursively solve for the children and only increment cnt if that edge is light. If at any stage the number of light edges is \gt log_2 n we return 0. If we reach a leaf then we return 1. See explanation for how to calculate this optimally.
Explanation
Since this is a tree, it is not much of a leap to guess that we need some sort of recursive solution. Suppose we start recursion from the root and solve recursively for children, what extra information do we need? And what do we do for each node? From the statement, it is clear that the action we must perform for each node is: choose each of the edges as a heavy edge and calculate the number of ways we can make a correct heavy light decomposition with that edge marked heavy. Sum up these ways for all the edges.
The only other constraint is that the number of light edges from root to any node must be \le log_2 n. So, besides the node index itself, we also need the number of light edges traversed so far. This gives us the state of our recursion: (x, cnt) where x is the current node index and cnt is the number of light edges traversed from the root to x.
Suppose the function solve does this, the final answer will be solve(1, 0).
How to calculate solve(x, cnt)?
x has say m children c_1, c_2 ... c_m. We make each of the m edges heavy in turn and increment cnt while calling recursively for all children except the one we made heavy. Now since the number of correct heavy light decomposition for one child is independent of the other children, we multiply the ways we get for each children. We sum this up for all possible choices of heavy edges to get the answer for solve(x, cnt). Since a state may repeat, we do memoization to avoid re calculations.
So,
solve(x, cnt) = \sum_{i = 1}^{m} (\prod_{j \ne i} solve(c_j, cnt + 1) * solve(c_i, cnt))
A straight forward implementation of this formula will TLE for subtask 3 as it will take O(m^2) = O(n^2) (Consider a tree where the root has n - 1 children).
To optimize let us see the inner product in detail.
We have \prod_{j \ne i} solve(c_j, cnt + 1) * solve(c_i, cnt).
We can split it as
\prod_{j = 1}^{i - 1} solve(c_j, cnt + 1) * \prod_{j = i + 1}^{m} solve(c_j, cnt + 1) * solve(c_i, cnt).
Why is this helpful? Well, we can see that we need two types of products: from left side till i - 1 and from right side till i + 1. We can precalculate this.
Let us define two arrays left and right.
left[i] = \prod_{j = 1}^{i} solve(c_j, cnt + 1)
right[i] = \prod_{j = i}^{m} solve(c_j, cnt + 1)
Now note that we can calculate this array in O(m):
left[0] = solve(c_0, cnt + 1)
left[j] = left[j - 1] * solve(c_j, cnt + 1)
and similarly for right.
Putting these two together:
solve(x, cnt) = \sum_{i = 1}^{m} (left[i - 1] * solve(c_i, cnt) * right[i + 1])
Remember that all calculations are modulo 19101995.
Complexity
How many possible states are there? There are n possible values for x and since we stop when cnt is log_2 n, the total number of states is O(n log n). Work done in each state is equal to the number of edges from that node. So, overall we can add them up and they will be equal to n - 1 - the number of edges. So, overall complexity is still O(nlogn)