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

  • Todo

  • print("Todo!")

All Problems

All Solutions