# Sum of Distances in Tree
**Difficulty:** HARD
[External](https://leetcode.com/problems/sum-of-distances-in-tree)
Canonical: https://scaleengineer.com/dsa/problems/sum-of-distances-in-tree
**Patterns:** [Dynamic Programming](https://scaleengineer.com/dsa/patterns/dynamic-programming)
**Algorithms:** [Depth-First Search](https://scaleengineer.com/algorithms/depth-first-search)
**Data structures:** Tree, Graph
**Companies:** [ByteDance](https://scaleengineer.com/companies/bytedance), [Media.net](https://scaleengineer.com/companies/media.net), [PhonePe](https://scaleengineer.com/companies/phonepe), [MathWorks](https://scaleengineer.com/companies/mathworks)
---
## Problem
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:**

![](https://assets.glich.co/dsa/sum-of-distances-in-tree/image0.jpg) 

**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:**

![](https://assets.glich.co/dsa/sum-of-distances-in-tree/image1.jpg) 

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

**Example 3:**

![](https://assets.glich.co/dsa/sum-of-distances-in-tree/image2.jpg) 

**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.

# Approaches
## Brute Force using BFS for each node
This approach is the most straightforward. For every single node in the tree, we can calculate the sum of its distances to all other nodes. To find the distance from a starting node `u` to all other nodes `v`, we can use a graph traversal algorithm like Breadth-First Search (BFS) or Depth-First Search (DFS). BFS is a natural choice here as it explores the tree layer by layer, directly giving the shortest distance from the source.
**Time:** O(N^2). For each of the N nodes, we perform a BFS traversal. A single BFS on a tree takes O(N + E) time, where E is the number of edges (N-1). So, one BFS is O(N). Repeating this for all N nodes results in a total time complexity of O(N * N) = O(N^2). · **Space:** O(N). The adjacency list requires O(N + E) = O(N) space. The queue and the `distances` array used in each BFS also require O(N) space.
**Pros:** Conceptually simple and easy to implement.; It's a direct translation of the problem statement.
**Cons:** Inefficient. The O(N^2) time complexity will result in a "Time Limit Exceeded" (TLE) for larger inputs (like N = 3 * 10^4).; It performs a lot of redundant calculations. The distance between two nodes `u` and `v` is calculated twice, once when starting from `u` and once when starting from `v`.
### Explanation
We iterate through each node `i` from `0` to `n-1`. For each `i`, we perform a BFS starting from it to compute the distances to all other nodes. A queue is used for the BFS, storing pairs of `(node, distance)`. We also need a `visited` array for each BFS run to avoid re-visiting nodes. During the BFS starting from `i`, we maintain a running sum. When we visit a node `j` at a distance `d`, we add `d` to this sum. After the BFS completes, the total sum is the answer for node `i`, so we store it in `answer[i]`. This process is repeated for all `n` nodes. The first step is to build an adjacency list representation of the tree from the input `edges` array to facilitate the traversal.
```java
import java.util.*;

class Solution {
    public int[] sumOfDistancesInTree(int n, int[][] edges) {
        if (n == 1) {
            return new int[]{0};
        }

        List<List<Integer>> adj = new ArrayList<>();
        for (int i = 0; i < n; i++) {
            adj.add(new ArrayList<>());
        }
        for (int[] edge : edges) {
            adj.get(edge[0]).add(edge[1]);
            adj.get(edge[1]).add(edge[0]);
        }

        int[] answer = new int[n];
        for (int i = 0; i < n; i++) {
            answer[i] = bfs(i, n, adj);
        }
        return answer;
    }

    private int bfs(int startNode, int n, List<List<Integer>> adj) {
        int[] dist = new int[n];
        Arrays.fill(dist, -1);
        Queue<Integer> queue = new LinkedList<>();

        queue.offer(startNode);
        dist[startNode] = 0;
        int totalDistance = 0;

        while (!queue.isEmpty()) {
            int u = queue.poll();
            totalDistance += dist[u];

            for (int v : adj.get(u)) {
                if (dist[v] == -1) {
                    dist[v] = dist[u] + 1;
                    queue.offer(v);
                }
            }
        }
        return totalDistance;
    }
}
```
### Algorithm
- 1. Create an adjacency list `adj` to represent the tree.
- 2. Initialize an `answer` array of size `n`.
- 3. For each node `i` from `0` to `n-1`:
    - a. Initialize a queue for BFS and add `i` to it.
    - b. Initialize a `distances` array of size `n` (e.g., filled with -1) and set `distances[i] = 0`.
    - c. Initialize a variable `currentSum = 0`.
    - d. While the queue is not empty:
        - i. Dequeue a node `u`.
        - ii. Add `distances[u]` to `currentSum`.
        - iii. For each neighbor `v` of `u`:
            - If `v` has not been visited (`distances[v] == -1`):
                - Set `distances[v] = distances[u] + 1`.
                - Enqueue `v`.
    - e. Store the result: `answer[i] = currentSum`.
- 4. Return the `answer` array.

## Optimal Two-Pass DFS Approach
A more efficient solution can be achieved by recognizing the relationship between the answers for adjacent nodes. Instead of re-calculating everything from scratch for each node, we can derive the answer for a node based on the answer of its parent. This dynamic programming approach on the tree requires two Depth-First Search (DFS) traversals.
**Time:** O(N). We perform two separate DFS traversals on the tree. Each traversal visits every node and edge exactly once. Building the adjacency list also takes O(N) time. Thus, the total time complexity is linear with respect to the number of nodes. · **Space:** O(N). The adjacency list requires O(N) space. The `count` and `ans` arrays also take O(N) space. The recursion stack for DFS can be up to O(N) in the worst case (for a skewed tree), leading to an overall O(N) space complexity.
**Pros:** Highly efficient with linear time complexity, making it suitable for large inputs.; It's the optimal way to solve this problem.
**Cons:** The logic is more complex than the brute-force approach.; It requires a good understanding of tree traversals and dynamic programming on trees.
### Explanation
The core idea is to first calculate the answer for an arbitrary root (say, node 0) and then use this result to efficiently calculate the answers for all other nodes. This is done in two passes.

**Pass 1: Post-order Traversal (Bottom-up)**
We perform a DFS starting from root 0. This traversal will compute two things for each node `u`:
- `count[u]`: The number of nodes in the subtree rooted at `u` (including `u`).
- `ans[u]`: The sum of distances from `u` to all other nodes *within its own subtree*.
The DFS explores children first. After returning from all children `v` of `u`, we can compute `count[u]` and `ans[u]` for the parent `u`.
- `count[u] = 1 + sum(count[v])` for all children `v`.
- `ans[u] = sum(ans[v] + count[v])` for all children `v`. The term `ans[v]` is the sum of distances from `v` to its descendants. To get the sum of distances from `u` to `v`'s descendants, we note that each path is one edge longer (the edge `u-v`), so we add `count[v]` (the number of such descendants).
After this pass, `ans[0]` holds the correct final answer for the root node, as its subtree is the entire tree.

**Pass 2: Pre-order Traversal (Top-down)**
We perform a second DFS, also starting from root 0. This traversal will use the parent's correct answer to calculate the child's correct answer. When we move from a parent `u` to a child `v`, the sum of distances changes.
- For the `count[v]` nodes in `v`'s subtree, their distance to the new "root" `v` becomes 1 shorter. This contributes a change of `-count[v]`.
- For the `n - count[v]` nodes outside `v`'s subtree, their distance to the new "root" `v` becomes 1 longer. This contributes a change of `+(n - count[v])`.
So, the update formula is: `ans[v] = ans[u] - count[v] + (n - count[v])`. We apply this formula as we traverse down the tree, and after this pass, the `ans` array is fully populated with the correct values.
```java
import java.util.*;

class Solution {
    private List<List<Integer>> adj;
    private int[] count;
    private int[] ans;
    private int n;

    public int[] sumOfDistancesInTree(int n, int[][] edges) {
        this.n = n;
        this.adj = new ArrayList<>();
        for (int i = 0; i < n; i++) {
            adj.add(new ArrayList<>());
        }
        for (int[] edge : edges) {
            adj.get(edge[0]).add(edge[1]);
            adj.get(edge[1]).add(edge[0]);
        }

        this.count = new int[n];
        this.ans = new int[n];

        // Pass 1: Post-order traversal to calculate subtree sizes and partial sums
        dfs1(0, -1);
        
        // Pass 2: Pre-order traversal to update sums for all nodes
        dfs2(0, -1);

        return ans;
    }

    // Post-order traversal: calculate count[u] and ans[u] for the subtree rooted at u
    private void dfs1(int u, int parent) {
        count[u] = 1;
        for (int v : adj.get(u)) {
            if (v == parent) continue;
            dfs1(v, u);
            count[u] += count[v];
            ans[u] += ans[v] + count[v];
        }
    }

    // Pre-order traversal: calculate the final answer for each node
    private void dfs2(int u, int parent) {
        for (int v : adj.get(u)) {
            if (v == parent) continue;
            // ans[v] is currently the sum of distances in its subtree.
            // We need to update it based on the parent's final answer.
            // ans[u] is the final answer for node u.
            // Move root from u to v:
            // - nodes in v's subtree get 1 closer: -count[v]
            // - nodes outside v's subtree get 1 farther: +(n - count[v])
            ans[v] = ans[u] - count[v] + (n - count[v]);
            dfs2(v, u);
        }
    }
}
```
### Algorithm
- 1. Build an adjacency list `adj` for the tree.
- 2. Initialize two arrays, `count` (for subtree sizes) and `ans` (for the final answer), both of size `n`.
- 3. **First Pass (post-order DFS, e.g., `dfs1(u, parent)`):**
    - a. Start DFS from an arbitrary root (e.g., node 0).
    - b. For a node `u`, recursively call `dfs1` for all its children `v`.
    - c. After the children's calls return, update `count[u]` and `ans[u]`:
        - `count[u] = 1 + Σ count[v]` for all children `v`.
        - `ans[u] = Σ (ans[v] + count[v])` for all children `v`.
- 4. **Second Pass (pre-order DFS, e.g., `dfs2(u, parent)`):**
    - a. Start DFS from the same root (node 0).
    - b. For a node `u`, before visiting its children, its `ans[u]` is already correct.
    - c. For each child `v` of `u`, calculate its final answer using the parent's answer:
        - `ans[v] = ans[u] - count[v] + (n - count[v])`.
    - d. Recursively call `dfs2` for the child `v`.
- 5. Return the `ans` array.

# Solutions
### Java

```java
class Solution {
private
  int n;
private
  int[] ans;
private
  int[] size;
private
  List<Integer>[] g;
public
  int[] sumOfDistancesInTree(int n, int[][] edges) {
    this.n = n;
    g = new List[n];
    ans = new int[n];
    size = new int[n];
    Arrays.setAll(g, k->new ArrayList<>());
    for (var e : edges) {
      int a = e[0], b = e[1];
      g[a].add(b);
      g[b].add(a);
    }
    dfs1(0, -1, 0);
    dfs2(0, -1, ans[0]);
    return ans;
  }
private
  void 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];
      }
    }
  }
private
  void 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]);
      }
    }
  }
}

```

### CPP

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

```

### Python

```python
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

```