834. Sum of Distances in Tree

📋 Đề Bài

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] = [a_i, b_i] indicates that there is an edge between nodes a_i and b_i in the tree.

Return an array answer of length n where answer[i] is the sum of the distances between the i^th 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 * 10⁴
edges.length == n - 1
edges[i].length == 2
0 <= a_i, b_i < n
a_i != b_i
The given input represents a valid tree.

🧠 Thuật Toán & Kỹ Thuật

DFS (Tìm kiếm theo chiều sâu)Hash Table (Bảng băm)Graph (Đồ thị)

⏱️ Thời gian O(V+E)

💾 Không gian O(V)

💻 Lời Giải

Python 0834-sum-of-distances-in-tree.py

class Solution:
    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:
        graph = defaultdict(list)

        for u, v in edges:
            graph[u].append(v)
            graph[v].append(u)

        cnt = [1] * n

        def dfs(u, parent):
            cntNode = 1
            for v in graph[u]:
                if v != parent:
                    cntNode += dfs(v, u)
            cnt[u] = cntNode
            return cntNode
        
        dfs(0, -1)
        ans = [0] * n

        def dfs2(u, parent, total):
            ans[u] = total
            for v in graph[u]:
                if v != parent:
                    dfs2(v, u, total + n - 2 * cnt[v])
    
        dfs2(0, -1, sum(cnt[1:]))
        return ans