PROBLEM LINK:
Problem Ideas: Tanuj Khattar, Arjun Arul, Malvika Raj Joshi
Author: Tanuj Khattar
Tester: Prashant Mahesh
Editorialist: Tanuj Khattar
DIFFICULTY:
HARD.
PREREQUISITES:
DP, Convex Hull Trick, DSU on Tree/Merging smaller to larger trick
PROBLEM:
Given a rooted unweighted tree where each node has some value, we need to partition the node of the tree into paths starting at some node and going towards the root, such that the following cost function is minimized.
[Cost(Partitioning\ of\ tree) = Sum(Cost\ of\ each\ path)]
[Cost(Path\ starting\ at\ node\ u)\ =\ \sum_{\forall\ w\ on\ the\ path}\left(C[u] \times (h  depth(w)) + C[u]^2 \right)  H[u] ]
EXPLANATION:
Analyzing the Cost Function:
Let cost(y,x) denote the cost of a path starting at node y and going up to node x such that x is an ancestor of y in the rooted tree. Then, as per the definition above,
Let d_{u} = h  depth(u). Hence, d_u for any node u represents its distance from the level at which the deepest node is located. Note that for any path (y,x), d_w's will be consecutive numbers, increasing as we go from y to x, where x is an ancestor of y in the tree.
Substituting in the above equation, we get:
On simplifying the summation, we get:
Where SUM(n) denotes the sum of first n natural numbers i.e. SUM(n) = \frac{n \times (n + 1)}{2}. On rearranging the above equation, we get:
Which looks like this:
where:
Note that cost(y,x) can be obtained by substituting d_{x} in the parabola A\cdot x^2 + B\cdot x + C where A,B,C depend only on parameters of y \left( C_{y},H_{y}, d_{y} \right) as shown above.
Convex Hull Trick for Parabolic Function:
Note that had the cost function above been linear instead of quadratic, we could have easily applied CHT to solve the linear version of the problem. But the given function is quadratic. However, note that both A and B depend only C_{y}. Due to this constraint, it is actually possible to apply CHT for the above cost function. Let us see why.
Let
Since A = C_{y} \geq 0 \forall y, Hence F(x) is an upwards opening parabola.

Claim1: F(x) is strictly increasing \forall x \geq 0
Proof: Notice that any upward opening parabola is strictly increasing for all x \geq \frac{B}{2\cdot A}. In this case, sign(B) = sign(A), hence sign( \frac{B}{2\cdot A}) = ve. Hence, F(x) is strictly increasing \forall x \geq 0**. 
Claim2: F_{y_{1}}(x) and F_{y_{2}}(x) can have at most one point of intersection \forall x \geq 0.
Proof: Let g(x) = F_{y_{1}}(x)  F_{y_{2}}(x) = A\prime \cdot x^2 + B\prime \cdot x + C\prime, where A\prime = A_{y_{1}}  A_{y_{2}}, B\prime = B_{y_{1}}  B_{y_{2}}, C\prime = C_{y_{1}}  C_{y_{2}}. Note that sign(A\prime) = sign(B\prime) due to the above mentioned fact that "For F(x) both A and B depend only on C_{y}". Therefore, g(x) can have atmost one nonnegative root. Hence proved.
Due to the above claims, we can say that the parabolas encountered due to the given cost function behave very similar to a line and hence we can apply Dynamic CHT on these parabolas instead of lines. Also, note that our query points for the CHT are always nonnegative. (d_{x} \geq 0 \forall x).
Solving the Array Version:
For the array version, it is very easy to come up with a straight forward \mathcal{O}(n^2) dp as follows:
where dp[x] denotes the min cost to paint all elements from x to n. Finally dp[1] is the answer.
Substituting the cost function, we get:
To improve the time complexity, we use CHT for parabolas to maintain the lower hull of parabolas corresponding to every y \geq x and then query for min at x = d_{x} in our CHT structure. This improves the time complexity to \mathcal{O}(n*log{}{n}) which is good.
Solving the Tree Version:
The idea remains the same but we need to be more careful with the details. An \mathcal{O}(n^2) dp for the tree would be:
In the above dp, we first add the dp values of children of x. Then, while iterating over all the y in subtree of x, we want to remove the dp value of child of x which is an ancestor of y, and then add the cost of path from y to x. Now, we also need to add the dp values of the remaining yellow and blue subtrees to get the final dp value of x if x belongs to path y>x. We take the min of this over all nodes y in subtree of x.
Notice that apart from the cost, every other value in the dp expression is independent of x. Hence we maintain a parabola for each node of the tree such that all the values in the above expression are incorporated in the constant of the parabola. Notice that while moving up by 1 edge, all the parabolas in that subtree are shifted up by the same constant due to the contribution of the new blue nodes and change in “child toward y” node for all nodes in the current subtree. Hence, we maintain a CHT structure at every node of the tree, which has parabolas corresponding to every node in it’s subtree. To get the CHT of a node x, we take the CHT of the heaviest child of x and then, iterate over parabolas in CHT of every other child and insert them one by one into the current CHT, with the updated constant based on the shifting factor of the two CHT’s. Since every parabola will be visited atmost \mathcal{O}(log{}{n}) times during the whole procedure, the complexity of this solution is \mathcal{O}(n*log^2{}{n}), which fits the TL with the given constraints.
The last part is a bit tricky so please try to work it out on pen and paper yourself. For more details, please have a look at the code below. I have tried to add comments wherever necessary. If you still have doubts, please ask in the comments below.
SOLUTIONS:
Author’s solution can be found here.