Chương trình nối rừng thành cây đơn với đường đi dài nhất nhỏ nhất trong Python

Cho rừng (tập hợp cây rời rạc) dạng danh sách kề. Thêm cạnh để nối thành 1 cây. Tìm độ dài đường đi dài nhất (diameter) nhỏ nhất có thể của cây kết quả.

Ví dụ rừng 2 cây → kết quả 4 (nối gốc 2 cây, diameter mới = max(đường đi qua cạnh mới, diameter cũ)).

Giải thuật (Tính bán kính cây + Nối tối ưu)

1. Với mỗi cây trong rừng: DFS tính diameter (đường đi dài nhất) và radius = ceil(diameter/2) (bán kính).
2. Tập hợp các radius của các cây: opt.
3. Nếu chỉ 1 cây → trả về diameter của nó.
4. Nối các cây: giảm radius lớn nhất đi 1 (do nối tại tâm), các radius còn lại +1 (khoảng cách đến tâm mới).
5. Kết quả = max(sum(2 radius lớn nhất sau điều chỉnh), max_diameter_cũ).

Ví dụ minh họa

import heapq, math

class Solution:
    def solve(self, graph):
        seen = set()
        
        def treeDiameter(node):
            self.dia = 0
            def dfs(u, p):
                seen.add(u)
                best = []
                for v in graph[u]:
                    if v != p:
                        heapq.heappush(best, dfs(v, u) + 1)
                        if len(best) > 2:
                            heapq.heappop(best)
                if not best:
                    return 0
                self.dia = max(self.dia, sum(best))
                return max(best)
            dfs(node, None)
            return self.dia
        
        opt = []
        max_dia = 0
        for i in range(len(graph)):
            if i in seen:
                continue
            dia = treeDiameter(i)
            max_dia = max(max_dia, dia)
            opt.append((dia + 1) // 2)  # radius = ceil(dia/2)
        
        if len(opt) <= 1:
            return max_dia
        
        # Nối cây: giảm radius max 1, các cái khác +1
        mx = max(opt)
        for i in range(len(opt)):
            if opt[i] == mx:
                opt[i] -= 1
                break
        opt = [r + 1 for r in opt]
        opt.sort(reverse=True)
        return max(opt[0] + opt[1], max_dia)

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

Đầu vào / Đầu ra

Input: danh sách kề rừng 2 cây

Output: 4