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 i
th 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); } } } }
-
Todo
-
print("Todo!")