CSES Solutions - Tree Distances II

You are given a tree consisting of n nodes. Your task is to determine for each node the sum of the distances from the node to all other nodes.

Example:

Input: edges = { { 1, 2 }, { 1, 3 }, { 3, 4 }, { 3, 5 } }
Output: 6 9 5 8 8

Input: edges = { { 1, 2 }, { 2, 3 }, { 3, 4 }, { 3, 5 } }
Output: 9 6 5 8 8

Approach:

Lets see the intuition into steps

1. Finding the sum of distances from a single node: If we root the tree at a particular node, we can use Depth-First Search (DFS) to find the depth of each other node. The sum of these depths gives us the sum of distances from the root node to all other nodes.
2. Finding the sum of distances from all nodes: Doing the above for each node would be too slow if n is large. So, we need a faster way. The key is to use the answer for one node to quickly find the answer for its neighbors.
3. Transitioning from one node to its neighbor: If we reroot the tree at a neighbor of the current root, the depths of all nodes in the new root’s subtree decrease by 1, and the depths of all nodes outside of its subtree increase by 1. This means the change in the sum of distances is exactly n - 2 * (size of new root's subtree).
4. Calculating the answer for all nodes: We can use DFS twice to find the answer for all nodes. The first DFS finds the answer for one node and the size of each node’s subtree. The second DFS computes the answer for all nodes using the formula from step 3.

Steps-by-step approach:

• DFS1 Function:
  • Performs a depth-first search (DFS) starting from node 1.
  • Calculates the answer for node 1 and initializes the size of each node's subtree.
  • For each child of the current node:
    • If the child is not the parent, it recursively calls dfs1() on the child, passing the child as the current node, the current node as the parent, and the depth incremented by 1.
    • It updates dp[node] by adding the size of the child's subtree.
• DFS2 Function:
  • Performs a depth-first search (DFS) starting from node 1.
  • For each child of the current node:
    • If the child is not the parent, it calculates the answer for the child using the formula provided in the comments.
    • It recursively calls dfs2() on the child, passing the child as the current node and the current node as the parent.
• Main Function:
  • Initializes the number of nodes n and the edges of the tree.
  • Builds the graph by adding edges to the adjacency list representation.
  • Calls dfs1() to compute initial values.
  • Calls dfs2() to calculate the final answer for each node.
  • Prints the answer for each node.

Below is the implementation of the above approach:

6 9 5 8 8 

Time Complexity: O(n)
Auxiliary Space: O(n)