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; }