Chương trình đếm số đường đi duy nhất chứa cạnh cho trước trong cây Python

Cho danh sách cạnh (u, v) biểu diễn một cây. Với mỗi cạnh, tìm số lượng đường đi duy nhất (unique paths) chứa cạnh đó. Kết quả trả về theo thứ tự cạnh đầu vào.

Ví dụ: edges = [[0,1],[0,2],[1,3],[1,4]] → kết quả [6, 4, 4, 4].

Giải thuật (DFS tính kích thước subtree)

1. Xây dựng danh sách kề adj từ các cạnh.
2. DFS từ gốc (0) để tính count[x] = số nút trong subtree gốc tại x.
3. Với mỗi cạnh (a, b), giả sử a là cha của b (hoặc ngược lại). Kích thước một bên là x = min(count[a], count[b]). Số đường đi qua cạnh đó = x * (N - x) với N = count[0].

Ví dụ minh họa

from collections import defaultdict

class Solution:
    def solve(self, edges):
        adj = defaultdict(list)
        for a, b in edges:
            adj[a].append(b)
            adj[b].append(a)
        
        count = defaultdict(int)
        
        def dfs(x, parent):
            count[x] = 1
            for nb in adj[x]:
                if nb == parent:
                    continue
                count[x] += dfs(nb, x)
            return count[x]
        
        dfs(0, -1)
        N = count[0]
        ans = []
        for a, b in edges:
            x = min(count[a], count[b])
            ans.append(x * (N - x))
        return ans

ob = Solution()
edges = [[0,1],[0,2],[1,3],[1,4]]
print(ob.solve(edges))  # Output: [6, 4, 4, 4]

Đầu vào / Đầu ra

Input: [[0,1],[0,2],[1,3],[1,4]]

Output: [6, 4, 4, 4]