Poga The War Stopper

At first, if stuck, we should try to find at least an O(n2) solution, even if it gives TLE.

With some effort or experience, it's not that hard to find a DP solution for that.
The recurrence relation is like,

dp[u] = min(dp[v] + a*(sub[u] - sub[v])2 + b) where v is any descendant of node u and sub[u] = sum of all node values in subtree rooted at u.

But of course, this will give a TLE. To optimize it, let's expand the relation first. That way, we get,

dp[u] = min(dp[v] + a*sub[u]2 -2*a*sub[u]*sub[v] + a*sub[v]2 + b)

= min(-2*a*sub[v]*sub[u] + a*sub[v]2 + dp[v]) + a*sub[u]2 + b

Those who are familiar with Convex Hull Trick optimization of DP can see some resemblance here.

The upward you go in the tree, the more -2*a*sub[v] decreases as a is non-negative. This is the m of y = mx + b format of CHT lines, and b = a*sub[v]2 + dp[v].

But, the problem isn't as straightforward as that. You can't maintain the list of lines of CHT for each node easily.

At this point, we have two approaches. To work with current recurrence relation, we can come up with an O(n*log2(n)) solution. We can do a DSU on Tree, also known as Sack.
For each node, we can rebuild its list of lines in CHT except for the big child's subtree by going to the big child subtree last.

But, we can do better by changing the recurrence relation a bit.
Instead of taking v from descendants, let's try taking v from ancestors. In that case, we'll actually be finding DP value of each edge instead of node. Because an ancestor can have many children, we need to cut the edge in the current path. This is like,

dp[u-v] = min(-2*a*sub[down]*sub[v] + a*sub[down]2 + dp[anc]) + a*sub[v]2 + b
This finds DP value up to edge u-v, where anc is an ancestor and down is the child of ancestor in the path from anc to u. To find answer to the original problem, we can add a dummy node to edge leaf node and find minimum answer out of all the new edges, but that's implementation details, can vary from person to person.

So, now we have a top-down approach. We can go down to each node and find DP value of each edge going down from the node. Now, when we find DP value of one edge and recursively go to the child connected by that edge, we are adding the new line created by this new DP value to the global list of lines of CHT. But, after all the work for this child is done, before going to the next child, we need to undo the change done to CHT.

To do that, we keep track of previous size of lines in CHT and the exact position where a line was replaced. Then, after recursion for this child is done, we can undo back the change at that exact position and restore back the size of lines too. This works because while adding a single line to CHT list, the list doesn't encounter big changes. Only the size is reduced, and exactly one position of the list has values replaced.

The idea is shown in the below CHT code.

But this still has a problem. In regular CHT, adding lines were amortized O(1) as each line is added and removed exactly once. But that is no longer the case here. Imagine a tree with chain of n/2 nodes, each adding a line. Then, from the last node of the chain, there n/2 more children, each adds a line that removes all previous lines. So, this makes our previous solution O(n^2).
The solution to this is doing a binary search in addLineToCHT to find the first position where a line has to be removed.
So, our final solution becomes:

Resources on CHT:
https://wcipeg.com/wiki/Convex_hull_trick
https://rezwanarefin01.github.io/posts/convex-hull-trick/ (This one is in Bangla, and pretty great stuff)