834. Sum of Distances in Tree

Problem

An undirected, connected tree with N nodes labelled 0...N-1 and N-1edges are given.

The ith edge connects nodes edges[i][0] and edges[i][1] together.

Return a list ans, where ans[i] is the sum of the distances between node iand all other nodes.

Example 1:

Input: N = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
Output: [8,12,6,10,10,10]
Explanation: 
Here is a diagram of the given tree:
  0
 / \
1   2
   /|\
  3 4 5
We can see that dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
equals 1 + 1 + 2 + 2 + 2 = 8.  Hence, answer[0] = 8, and so on.

Note: 1 <= N <= 10000

Problem Summary

Given an undirected, connected tree. The tree has N nodes labeled 0...N-1 and N-1 edges. The ith edge connects nodes edges[i][0] and edges[i][1] . Return a list ans representing the sum of distances between node i and all other nodes.

Note: 1 <= N <= 10000

Solution Approach

Given N nodes and the relationships of some edges between these nodes. The requirement is to find the path sum from x as the root node to all nodes, for each x.
Although this problem describes finding paths in a tree, it can be treated entirely as a graph, because it is not a binary tree but a multiway tree. The solution idea is to first perform one DFS to find the path sums from 0 as the root node to each node (it does not have to use 0 as the node; any node can be chosen as the start). The second DFS computes the path sums after converting from root node 0 to each of the other nodes. Since the path sum computed the first time with node 0 as the root is correct, the path sum for other nodes as root nodes can be obtained correctly by just converting it. After 2 DFS traversals, we can get the path sum from every node as the root to all nodes.
How do we convert the path sum from root node 0 to all other nodes into the path sum from other nodes as root nodes to all nodes? Changing from node 0 to node x only requires adding and subtracting based on the path sum from 0 to all nodes. What increases is the paths from node x to the nodes outside all subtrees rooted at x; however many such nodes there are, that many paths are added. What decreases is the paths from 0 to all subtree nodes rooted at x, including from 0 to root node x; however many nodes there are, that many paths are subtracted. Therefore, in the first DFS, we need to calculate the total number of nodes in the subtree rooted at each node (including itself), so that in the second DFS we can directly use it for the conversion. See the code for the specific implementation details.

Code


package leetcode

func sumOfDistancesInTree(N int, edges [][]int) []int {
	// count[i] stores the total number of all subtree nodes and the root node when i is the root node
	tree, visited, count, res := make([][]int, N), make([]bool, N), make([]int, N), make([]int, N)
	for _, e := range edges {
		i, j := e[0], e[1]
		tree[i] = append(tree[i], j)
		tree[j] = append(tree[j], i)
	}
	deepFirstSearch(0, visited, count, res, tree)
	// Reset the visited state, then perform DFS again
	visited = make([]bool, N)
	// Before entering the second DFS, only res[0] stores the correct value, because the first DFS computed the sum of all paths with 0 as the root node
	// The purpose of the second DFS is to convert the path sum with 0 as the root node into the path sum with n as the root node
	deepSecondSearch(0, visited, count, res, tree)

	return res
}

func deepFirstSearch(root int, visited []bool, count, res []int, tree [][]int) {
	visited[root] = true
	for _, n := range tree[root] {
		if visited[n] {
			continue
		}
		deepFirstSearch(n, visited, count, res, tree)
		count[root] += count[n]
		// Sum of all paths from the root node to n = res[n], the path sum from n as the root node to all subtrees, + the number of each node in count[n] from root to count[n] (the root node and each node rooted at n each add one path)
		// The root node and each node rooted at n each add one path = when n is the root node, the total number of subtree nodes and the root node, i.e. count[n]
		res[root] += res[n] + count[n]
	}
	count[root]++
}

// Starting from root, set each child node of the root node as the new root node in turn
func deepSecondSearch(root int, visited []bool, count, res []int, tree [][]int) {
	N := len(visited)
	visited[root] = true
	for _, n := range tree[root] {
		if visited[n] {
			continue
		}
		// After the root node changes from root to n
		// res[root] stores the total path length from root as the root node to all nodes
		// 1. The increased path length from root to node n = the root node and each node rooted at n each add one path = when n is the root node, the total number of subtree nodes and the root node, i.e. count[n]
		// 2. The increased path length from n to nodes outside all subtree nodes rooted at n = node n and each node in the non-n-rooted subtree each add one path = N - count[n]
		// Therefore, to transfer the root node from root to n, the paths that need to be increased are calculated in the second step above 👆, and the paths that need to be decreased are calculated in the first step above 👆
		res[n] = res[root] + (N - count[n]) - count[n]
		deepSecondSearch(n, visited, count, res, tree)
	}
}