Problem Link
Setter: Hasan Jaddouh
Tester: Amr Mahmoud
Editorialist: Hasan Jaddouh
Difficulty
HARD
Prerequisites
Dynamic programming on trees, combinatorics
Problem
Given a tree, for each node i make it the root of the tree then count the number of subsets of leaves that it’s possible to be visited next to each other when doing dfs from the root (ignoring nodes that are not leaves)
Explanation
solution for a single root
Let’s root the tree at node 1 and count the number of good subsets of leaves, we will use dynamic programming technique.
dp_{i,0} means the number of good subsets in subtree of node i excluding empty subset and subset that consists of all leaves in the subtree (this is to avoid repetition in counting)
dp_{i,1} means the number of good subsets in subtree of node i that also appear as prefix in one of the dfs sequences (let’s call it good prefix subset), again excluding empty subset and subset that consists of all leaves in the subtree
Now suppose we have already calculated dp values of children of node i, how can we use them to calculate the dp values for node i itself?
first thing we notice that all those subsets of leaves that their LCA is a decedent of node i are already counted in one of children of node i so in the end we only add do values of the children to node i, so we only care about subsets of leaves that their LCA is node i
note that when we do dfs from node i we can select the order of visiting the children, but once we visited one child of node i we can’t start visiting another child of node i before we visit all leaves in subtree of the first child, so for a subset of leaves to be good then for each child of node i either all leaves in its subtree are included in the subset or none of them are included, except for at most two children because they can be a prefix and suffix in the sequence of visiting the good subset.
From what is said above we can have some insight about how to calculate the dp values, suppose C is the number of children of node i, we have the following cases of good subsets
-
for each non-empty subset S of children of node i, if we take all leaves that are in subtree of one of the children in S then it will be good subset of leaves, this also can be “good prefix subset” (we defined this above) so we add the count of this case to both dp_{i,0} and dp_{i,1}. the number of subsets of children of node i is 2^{C} we should subtract 2 because we don’t want to include the empty subset and we don’t want to include the full subset (remember we excluded it in our definition of dp_{i,0} and dp_{i,1} ).
-
for each child u of node i and a subset S of other children, if we take all leaves that are in subtree of one of the children in subset S in addition to some but not all leaves in subtree of u such that those leaves make good prefix subset then all those leaves will form a good subset of leaves and it also can be a prefix so we add the count of this case to both dp_{i,0} and dp_{i,1}. to count the number of subsets of this case we should add values of dp_{u,1} where u is child of i and then multiply it by 2^{C-1} -1, why C-1 instead of C? because we can’t include u itself, why -1 in the end? because we don’t want to include the empty subset of children otherwise this will make the resulting subset of leaves have descendant of node i as their LCA and lead to repetition of counting.
-
for each pair of different children u, v of node i and a subset S of other children, if we take all leaves that are in subtree of one of the children in S and some but not all leaves in subtree of u such that those leaves make good prefix subset and some but not all leaves in subtree of v such that those leaves make good prefix subset then all those leaves will form a good subset of leaves, but they can’t be good prefix subset so we only add this count to dp_{i,0}, to count the number of subsets in this case we should add all values of dp_{v,1} * dp_{u,1} where u and v are different children of node i after that we multiply the result by 2^{C-2} because this is the number of subsets of other children of node i.
in the end the answer is dp_{1,0} + 1, we add 1 to include the subset of all leaves again.
This solution only requires a single dfs to calculate all dp values so its complexity is O(n)
solution for all roots
Let’s start by rooting the tree at arbitrary node (say node 1) and apply the previous solution on it. Now consider a node that is not a root, it have a number of children (say C) and one parent, we calculated the dp values of this node from its C children, what would have been different if this node was the root? simply all its children are still the same but the parent will become a child of this node making this node have C + 1 children, so all we have to do is to calculate for each node what would have the dp values of its parent be if this parent was a child instead, and then calculate the dp values of all nodes again taking into account the new values from the parent, how to this exactly?
first of all, we do dfs from the the node we rooted the tree at (say node 1) and calculate dp values for all subtrees exactly the same in “solution for a single root” paragraph, after that we should do another dfs from the same root, but this time we should allow dfs function to have 2 more parameters that we will use for passing the dp values of the parent to its children, inside the dfs function first we should calculate again the dp values of the current node including the dp values that we got from the parameters, after that we will get the final answer for this node,note that calculations are done exactly the same way we did in first paragraph, after that we should calculate the parameters to pass to the children of current node, note that those parameters will be different for each child, and it’s also done in same way but only considering only dp values of all other children in addition to the dp values of the parent.
using this method we only need to do two dfs traversal on the tree, thus the whole complexity is O(n)
Time Complexity
O(N), per test case.