PROBLEM LINK:
Author: Zi Song Yeoh
DIFFICULTY:
EASY
PREREQUISITES:
DFS, TREES
PROBLEM:
Given a weighted tree, find the sum of heights of all nodes if each vertex is chosen as root.
EXPLANATION:
First, let’s solve the problem when for only one root. This is trivial, as we can just dfs once to find the height of each node and sum them up. Repeating this for all nodes gives an O(n^{2}) algorithm, which will work for Subtask 1.
For subtask 2, we need something better. The trick is that the answer won’t change much if we change the root to its neighbour. Formally, suppose r is the root and u is one of its neighbours. Let w be the weight of the edge connecting them. Then, if we shift the root from r to u, the total sum of heights will increase by (V - U)w, where U is the size of subtree rooted at u and V = n - U. Thus, we can calculate the answer for all nodes with a single dfs in O(n) time.
AUTHOR’S AND TESTER’S SOLUTIONS:
Author’s solution can be found here.
Tester’s solution can be found here.