834. Sum of Distances in Tree

[F4NT0]

There is an undirected connected tree with n nodes labeled from 0 to n - 1 and n - 1 edges.

You are given the integer n and the array edges where edges[i] = [ai, bi] indicates that there is an edge between nodes ai and bi in the tree.

Return an array answer of length n where answer[i] is the sum of the distances between the ith node in the tree 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: The tree is shown above.
• 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.

Example 2

• Input: n = 1, edges = []
• Output: [0]

Example 3

• Input: n = 2, edges = [[1,0]]
• Output: [1,1]

Constraints:

• 1 <= n <= 3 * 104
• edges.length == n - 1
• edges[i].length == 2
• 0 <= ai, bi < n
• ai != bi
• The given input represents a valid tree.

[F4NT0]

using System;
using System.Collections.Generic;

public class Solution {
    public int[] SumOfDistancesInTree(int n, int[][] edges) {
        int[] ans = new int[n];
        int[] count = new int[n];
        HashSet<int>[] tree = new HashSet<int>[n];

        for (int i = 0; i < n; ++i) {
            tree[i] = new HashSet<int>();
            count[i] = 1;
        }

        foreach (int[] edge in edges) {
            int u = edge[0];
            int v = edge[1];
            tree[u].Add(v);
            tree[v].Add(u);
        }

        Postorder(tree, 0, -1, count, ans);
        Preorder(tree, 0, -1, count, ans);
        return ans;
    }

    private void Postorder(HashSet<int>[] tree, int node, int parent, int[] count, int[] ans) {
        foreach (int child in tree[node]) {
            if (child == parent)
                continue;
            Postorder(tree, child, node, count, ans);
            count[node] += count[child];
            ans[node] += ans[child] + count[child];
        }
    }

    private void Preorder(HashSet<int>[] tree, int node, int parent, int[] count, int[] ans) {
        foreach (int child in tree[node]) {
            if (child == parent)
                continue;
            ans[child] = ans[node] - count[child] + (tree.Length - count[child]);
            Preorder(tree, child, node, count, ans);
        }
    }
}