### PROBLEM LINK:

Practice

Contest: Division 1

Contest: Division 2

**Setter:** Praveen Dhinwa

**Tester:** Teja Vardhan Reddy

**Editorialist:** Taranpreet Singh

### DIFFICULTY:

Medium

### PREREQUISITES:

Trees, Dynammic Programming.

### PROBLEM:

Given a Tree with N nodes and N markers, each being either 0, 1 and 2, we need to assign markers to each node in such a way to minimize, say ans = abs(m(u_i)-m(v_i)) for all edges (u_i, v_i) in tree if m(x) represent mark assigned to node x.

### QUICK EXPLANATION

- Dynamic Programming on tree, with state (u, i, j, cur), calling it outer DP representing the minimum ans in subtree of node x, if i markers labeled 0, j markers labeled 1 were used, with label cur assigned to node u.
- Assuming we have it calculated for all children, to merge results of all children, we use another DP table with state (ch, x, y, contains) representing ans when out of subtrees of first ch children, x children are marked 0, y children are marked 1 and contains is a bitmask with 3 bits, ith representing whether any of the first ch children have marker i assigned to it.
- Using internal DP for calculating the convolution of results of children so as to efficiently calculate outer DP state for current node efficiently.
- Final answer is represented as the min(0, cnt0, cnt1, i) where i is the marker assigned to node 0 and cnt0 and cnt1 are the Numbers of markers assigned label 0 and 1 respectively.

### EXPLANATION

So, Another Tree DP problem. This editorial will have a bit on how to actually solve Tree DP problem, so might be a bit lengthy.

Firstly, define the minimum difference in the subtree of node u as the maximum of absolute difference of markers assigned to any pair of nodes connected by an edge.

First of all, see, there are N markers and N nodes, so each will be assigned one of the markers. Easy, right?

Let’s root the tree at any node (Assuming rooted at 1) and C_i denote Number of markers marked i given.

Suppose considering the subtree of a node only. One marker will be assigned to the current node, while others will be assigned in the subtree of children. So, **the minimum difference we get is the maximum of the absolute difference between the current marker and its direct children, and the minimum difference in the subtree of any child.**

We see, that if we know the minimum difference calculated for its children and the markers assigned to children and the marker assigned to the current child, we can know the minimum difference.

Also, we need to use C_0 markers marked 0, C_1 markers marked 1 and C_2 markers marked 2, so, we also need to maintain the count of markers of each label.

Usually in tree problems where we can compute results for a node using the results of its children nodes and some additional information, then we can use Dynamic Programming to calculate results sequentially from leaf to root direction. We recursively compute information for the subtree of a node using information of subtrees of its children, then this computed info is used to compute more info on its parent and so on, till we get the info of subtree of the root, which is actually the information about the tree we seek.

The next question arising in such problems is the information we need to store about subtree rooted at each node so that we can calculate the same information for the parent node. Most of the time, we can get this information from the problem statement itself. For this problem, see, we need to answer Minimum difference in the subtree of root using given number of markers of each label. This hints us to include the number of markers of each label in the subtree of a node.

So, Now our Dynamic Programming state become (u, c0, c1, c2). But, we can reduce one state out of this by noticing the fact that we always have c0+c1+c2 = sub[u] where sub[u] is the size of the subtree rooted at u which always remain constant. This way, we have Dynamic state (u, c0, c1).

Now, reread the bold line and see what information we still need so as to calculate information for a node using information about its children. We see We also need information about the markers assigned to direct children as well as the marker assigned to the current child. Basically, we need the information whether any child is assigned label 0, label 1 and label 2 respectively. So, we get another required information which should be altogether calculated.

So, Our Dynamic Programming state become (u, c0, c1, cur) representing the minimum absolute difference in subtree of u such that c0 markers labeled 0, c1 markers labeled 1 and (sub[u]-c0-c1) markers labeled 2 are used in subtree of u such that the label assigned to node u is cur.

We can see, that we can compute information about parent using this information of all children.

Let’s move toward Base cases after assigning all states as some maximum value (Any value \geq 2 shall work) since each state will store a minimum value.

The base case is when node u is a leaf. We can assign any marker to this position. If assigned marker 0, c0 become 1 while c1 is 0 and minimum difference in this subtree is 0. So our state (u, 1, 0, 0) gets assigned value 0. We can work out similarly by assigning markers labeled 1 and 2, and assign states (u, 0, 1,1) and (u, 0,0,2) the value 0 since we don’t have any edge in the subtree of this node, hence the minimum absolute difference is zero.

Moving toward Dynamic Programming Transition.

See, that when calculating the answer for a state (u, c0, c1, cur) one marker labeled cur is used by the current node. Now, the remaining labels are used by children nodes in their subtrees. But we do not know which child has stored how many markers of each label.

For this, we shall use another Dynamic Programming table, to merge information of all children of a node. This DP will be calculated for every node.

We will try to merge information about children one by one from left to right sequentially.

Once again, see what information we need combining about first x children of u such that if we combine this information with information of our previous DP for x+1 child, We can get the same set of information for first x+1 children of u.

We can apply similar reasoning as above to notice that our state becomes (x, c0, c1) where x is number of children from left, about whom this info is calculated.

Interesting question is, earlier we had only one marker which is assigned to root of the subtree. But now, labels assigned to all **direct children** matter, so as to compute the minimum difference. So, we also need a state variable to determine to know, reflecting labels of first x children. Here, the fact that there are only three markers, come to our help. We can just make a bitmask of 3 bits, first bit from right, if this is an on bit, means first x children have at least 1 children with label 0 assigned and same for label 1 and 2.

So, Our DP state now becomes (x, c0, c1, mask) which represent the Minimum absolute difference in first x children of Node u such that they together use c0 markers labeled 0 and c1 labels marked 1. Note that the mask values will be in the range [0, 2^3-1] as we can represent all combinations of off and on bits in this range.

Now, Time for Transition in this DP table.

We see, that DP(x+1, a+p, b+q, mask|(cur bit)) = min(max(DP(x, a, b, mask), ANS(v, p, q, cur))) for all valid combinations of values of (a,b,p mask,p,q,cur) where ANS table is our main DP. DP(x, a, b, mask) is a state of internal DP while ANS(v,p, q, cur) is our main DP where v is x+1 child of u.

The transition is not hard to follow where we just take the sum of counts of labels used, except for the mask part, explaining now.

See, first x children, if no child has label cur, cur bit of mask will be an off bit. But when this state gets merged with a state having cur label, cur bit has to be turned on, since now, first x+1 children have at least 1 child with label cur.

This way, we have information about all children calculated. Time to calculate our ANS state for the current node. We see, that the number of nodes of type c0, and c1 will remain same as that when combined all children, except for the label assigned to the current node, for which, it will be increased by 1.

So, our DP transition becomes ANS(u, c0+(cur==0), c1+(cur==1), cur) = min(max(f(mask, cur), DP(X, c0, c1, mask))) for all valid combinations of c0, c1 mask and cur represent the label we are trying to assign to current node, and f(mask, cur) calculate the minimum absolute difference of edges between node u and its direct children.

This way, we can compute the whole ANS table, thus getting the respective values for root.

Now, we can easily find the minimum difference using c0 labels marked 0 and c1 labels marked 1 for all of the three labels we can assign to the current node, take the minimum. We have found the answer.

**Implementation**

Above implementation has a time complexity of O(N^3) but with a high constant factor which may lead to TLE. There are two important optimizations required.

First thing, we do not need to compute all states for all nodes. We are given how many markers of each label we are gonna use in the whole tree, so, it doesn’t make sense to calculate any state which involves the use of more markers of that label.

Secondly, We should avoid iterating over states, for which answer can be found beforehand. For example, see the transition of the DP table. In that, suppose we have ANS(x, c0, c1, mask) as MX value, the final value of the expression will be MX since we are taking maximum. This state is useless, so we directly skip over it.

Other useless states are the one which contradicts themselves. For example, state DP(u, 0,0,0) is an invalid state, since if node u is assigned 0, c0 should be at least 1. Another example, ANS(x, 0, 0, 011_2) since the mask says there is at least 1 node with label 1 while c1 is zero. We need to skip over these states too.

**Note About Internal DP and Convolution: Proceed at your own risk**

The internal DP is nothing but the multiplication convolution of trivariate polynomials if we ignore masks. Though doubtful if convolution can be actually applied to solve this problem. Let the discussion begin in comments.

**Related Problems**

This problem uses the same trick combined with combinatorics and bitwise operations.

Another problem here uses the same trick, combined with Probabilities and Expectation.

### Time Complexity Analysis

When we are running the internal DP, we perform sum*sub[v] where the sum is the number of nodes in first x children. We can split this sum to parts as t prove that we can write the number of operations in terms of the number of pairs of nodes, leading to O(N^2) time with merging taking O(N) time, leading to overall Time complexity to O(N^3).

### AUTHOR’S AND TESTER’S SOLUTIONS:

Setter’s solution

Tester’s solution

Editorialist’s solution

Feel free to Share your approach, If it differs. Suggestions are always welcomed.