Welcome to Subscribe On Youtube

Formatted question description: https://leetcode.ca/all/834.html

834. Sum of Distances in Tree

Level

Hard

Description

An undirected, connected tree with N nodes labelled 0...N-1 and N-1 edges 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 i and 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

Solution

Create two arrays counts and sums of length N, where counts[i] represents the number of nodes in the subtree that has root i and sums is the result array. Initialize all elements in counts to 1.

Do depth first search starting from node 0 for two times.

For the first time, calculate the array counts and each root’s value of sums. For each node i that is not a leaf, counts[i] is the sum of all elements counts[j] plus 1 and sums[i] is the sum of all elements sums[j] and counts[j], where j is a child of i.

For the second time, calculate each child’s value of sums. For each node i that is not 0, if its parent is j, then sums[i] = sums[j] + counts[i] + N - counts[i].

  • class Solution {
        public int[] sumOfDistancesInTree(int N, int[][] edges) {
            Map<Integer, Set<Integer>> edgesMap = new HashMap<Integer, Set<Integer>>();
            for (int[] edge : edges) {
                int node1 = edge[0], node2 = edge[1];
                Set<Integer> set1 = edgesMap.getOrDefault(node1, new HashSet<Integer>());
                Set<Integer> set2 = edgesMap.getOrDefault(node2, new HashSet<Integer>());
                set1.add(node2);
                set2.add(node1);
                edgesMap.put(node1, set1);
                edgesMap.put(node2, set2);
            }
            int[] counts = new int[N];
            Arrays.fill(counts, 1);
            int[] sums = new int[N];
            depthFirstSearch1(0, -1, edgesMap, counts, sums);
            depthFirstSearch2(0, -1, edgesMap, counts, sums, N);
            return sums;
        }
    
        public void depthFirstSearch1(int node, int parent, Map<Integer, Set<Integer>> edgesMap, int[] counts, int[] sums) {
            Set<Integer> children = edgesMap.getOrDefault(node, new HashSet<Integer>());
            for (int child : children) {
                if (child != parent) {
                    depthFirstSearch1(child, node, edgesMap, counts, sums);
                    counts[node] += counts[child];
                    sums[node] += sums[child] + counts[child];
                }
            }
        }
    
        public void depthFirstSearch2(int node, int parent, Map<Integer, Set<Integer>> edgesMap, int[] counts, int[] sums, int length) {
            Set<Integer> children = edgesMap.getOrDefault(node, new HashSet<Integer>());
            for (int child : children) {
                if (child != parent) {
                    sums[child] = sums[node] - counts[child] + length - counts[child];
                    depthFirstSearch2(child, node, edgesMap, counts, sums, length);
                }
            }
        }
    }
    
  • class Solution {
    public:
        vector<int> sumOfDistancesInTree(int n, vector<vector<int>>& edges) {
            vector<vector<int>> g(n);
            for (auto& e : edges) {
                int a = e[0], b = e[1];
                g[a].push_back(b);
                g[b].push_back(a);
            }
            vector<int> ans(n);
            vector<int> size(n);
    
            function<void(int, int, int)> dfs1 = [&](int i, int fa, int d) {
                ans[0] += d;
                size[i] = 1;
                for (int& j : g[i]) {
                    if (j != fa) {
                        dfs1(j, i, d + 1);
                        size[i] += size[j];
                    }
                }
            };
    
            function<void(int, int, int)> dfs2 = [&](int i, int fa, int t) {
                ans[i] = t;
                for (int& j : g[i]) {
                    if (j != fa) {
                        dfs2(j, i, t - size[j] + n - size[j]);
                    }
                }
            };
    
            dfs1(0, -1, 0);
            dfs2(0, -1, ans[0]);
            return ans;
        }
    };
    
  • class Solution:
        def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:
            def dfs1(i: int, fa: int, d: int):
                ans[0] += d
                size[i] = 1
                for j in g[i]:
                    if j != fa:
                        dfs1(j, i, d + 1)
                        size[i] += size[j]
    
            def dfs2(i: int, fa: int, t: int):
                ans[i] = t
                for j in g[i]:
                    if j != fa:
                        dfs2(j, i, t - size[j] + n - size[j])
    
            g = defaultdict(list)
            for a, b in edges:
                g[a].append(b)
                g[b].append(a)
    
            ans = [0] * n
            size = [0] * n
            dfs1(0, -1, 0)
            dfs2(0, -1, ans[0])
            return ans
    
    
  • func sumOfDistancesInTree(n int, edges [][]int) []int {
    	g := make([][]int, n)
    	for _, e := range edges {
    		a, b := e[0], e[1]
    		g[a] = append(g[a], b)
    		g[b] = append(g[b], a)
    	}
    	ans := make([]int, n)
    	size := make([]int, n)
    	var dfs1 func(i, fa, d int)
    	dfs1 = func(i, fa, d int) {
    		ans[0] += d
    		size[i] = 1
    		for _, j := range g[i] {
    			if j != fa {
    				dfs1(j, i, d+1)
    				size[i] += size[j]
    			}
    		}
    	}
    	var dfs2 func(i, fa, t int)
    	dfs2 = func(i, fa, t int) {
    		ans[i] = t
    		for _, j := range g[i] {
    			if j != fa {
    				dfs2(j, i, t-size[j]+n-size[j])
    			}
    		}
    	}
    	dfs1(0, -1, 0)
    	dfs2(0, -1, ans[0])
    	return ans
    }
    
  • function sumOfDistancesInTree(n: number, edges: number[][]): number[] {
        const g: number[][] = Array.from({ length: n }, () => []);
        for (const [a, b] of edges) {
            g[a].push(b);
            g[b].push(a);
        }
        const ans: number[] = new Array(n).fill(0);
        const size: number[] = new Array(n).fill(0);
        const dfs1 = (i: number, fa: number, d: number) => {
            ans[0] += d;
            size[i] = 1;
            for (const j of g[i]) {
                if (j !== fa) {
                    dfs1(j, i, d + 1);
                    size[i] += size[j];
                }
            }
        };
        const dfs2 = (i: number, fa: number, t: number) => {
            ans[i] = t;
            for (const j of g[i]) {
                if (j != fa) {
                    dfs2(j, i, t - size[j] + n - size[j]);
                }
            }
        };
        dfs1(0, -1, 0);
        dfs2(0, -1, ans[0]);
        return ans;
    }
    
    

All Problems

All Solutions