834. Sum of Distances in Tree LeetCode Solution

Last updated: February 20, 2025

In this guide, you will get 834. Sum of Distances in Tree LeetCode Solution with the best time and space complexity. The solution to Sum of Distances in Tree problem is provided in various programming languages like C++, Java, and Python. This will be helpful for you if you are preparing for placements, hackathons, interviews, or practice purposes. The solutions provided here are very easy to follow and include detailed explanations.

Problem Statement
Complexity Analysis
Sum of Distances in Tree solution in C++
Sum of Distances in Tree solution in Java
Sum of Distances in Tree solution in Python
Additional Resources

834. Sum of Distances in Tree LeetCode Solution image

Problem Statement of Sum of Distances in Tree

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.

Complexity Analysis

Time Complexity: O(n)
Space Complexity: O(n)

834. Sum of Distances in Tree LeetCode Solution in C++

class Solution {
 public:
  vector<int> sumOfDistancesInTree(int n, vector<vector<int>>& edges) {
    vector<int> ans(n);
    vector<int> count(n, 1);
    vector<unordered_set<int>> tree(n);

    for (const vector<int>& edge : edges) {
      const int u = edge[0];
      const int v = edge[1];
      tree[u].insert(v);
      tree[v].insert(u);
    }

    postorder(tree, 0, /*prev=*/-1, count, ans);
    preorder(tree, 0, /*prev=*/-1, count, ans);
    return ans;
  }

 private:
  void postorder(const vector<unordered_set<int>>& tree, int u, int prev,
                 vector<int>& count, vector<int>& ans) {
    for (const int v : tree[u]) {
      if (v == prev)
        continue;
      postorder(tree, v, u, count, ans);
      count[u] += count[v];
      ans[u] += ans[v] + count[v];
    }
  }

  void preorder(const vector<unordered_set<int>>& tree, int u, int prev,
                vector<int>& count, vector<int>& ans) {
    for (const int v : tree[u]) {
      if (v == prev)
        continue;
      // count[v] us are 1 step closer from v than prev.
      // (n - count[v]) us are 1 step farther from v than prev.
      ans[v] = ans[u] - count[v] + (tree.size() - count[v]);
      preorder(tree, v, u, count, ans);
    }
  }
};
/* code provided by PROGIEZ */

834. Sum of Distances in Tree LeetCode Solution in Java

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

    Arrays.fill(count, 1);

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

    for (int[] edge : edges) {
      final int u = edge[0];
      final 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(Set<Integer>[] tree, int u, int prev, int[] count, int[] ans) {
    for (final int v : tree[u]) {
      if (v == prev)
        continue;
      postorder(tree, v, u, count, ans);
      count[u] += count[v];
      ans[u] += ans[v] + count[v];
    }
  }

  private void preorder(Set<Integer>[] tree, int u, int prev, int[] count, int[] ans) {
    for (final int v : tree[u]) {
      if (v == prev)
        continue;
      // count[v] us are 1 step closer from v than prev.
      // (n - count[v]) us are 1 step farther from v than prev.
      ans[v] = ans[u] - count[v] + (tree.length - count[v]);
      preorder(tree, v, u, count, ans);
    }
  }
}
// code provided by PROGIEZ

834. Sum of Distances in Tree LeetCode Solution in Python

class Solution:
  def sumOfDistancesInTree(self, n: int, edges: list[list[int]]) -> list[int]:
    ans = [0] * n
    count = [1] * n
    tree = [set() for _ in range(n)]

    for u, v in edges:
      tree[u].add(v)
      tree[v].add(u)

    def postorder(u: int, prev: int) -> None:
      for v in tree[u]:
        if v == prev:
          continue
        postorder(v, u)
        count[u] += count[v]
        ans[u] += ans[v] + count[v]

    def preorder(u: int, prev: int) -> None:
      for v in tree[u]:
        if v == prev:
          continue
        # count[v] us are 1 step closer from v than prev.
        # (n - count[v]) us are 1 step farther from v than prev.
        ans[v] = ans[u] - count[v] + (n - count[v])
        preorder(v, u)

    postorder(0, -1)
    preorder(0, -1)
    return ans
 # code by PROGIEZ

Additional Resources

Happy Coding! Keep following PROGIEZ for more updates and solutions.

Table of Contents