Chương trình tìm thời gian lan truyền tin nhắn đến tất cả nút trong mạng lưới trong Python

Cho số nút n (đánh số 0..n) và danh sách cạnh edges dạng (a, b, t) (vô hướng, trọng số t là thời gian truyền). Tin nhắn bắt đầu từ nút 0, mỗi nút nhận được sẽ lan truyền ngay đến các nút kề. Tìm thời gian sớm nhất mà tất cả nút nhận được tin nhắn.

Ví dụ: n=3, edges=[[0,1,3],[1,2,4],[2,3,2]] → nút 3 nhận tin sau 3+4+2=9.

Giải thuật (Dijkstra - Tìm đường đi ngắn nhất từ nguồn đơn)

• Xây dựng đồ thị danh sách kề.
• Dùng hàng đợi ưu tiên (min-heap) lưu (total_time, node).
• BFS Dijkstra: pop nút có thời gian nhỏ nhất, nếu chưa thăm thì đánh dấu đã thăm. Nếu đã thăm hết n+1 nút → trả về thời gian hiện tại.
• Đẩy các nút kề chưa thăm với thời gian cộng dồn.

Mã nguồn Python

import heapq
from collections import defaultdict

class Solution:
    def solve(self, n, edges):
        graph = self.build_graph(edges)
        visited = set()
        heap = [(0, 0)]
        while heap:
            current_total_time, node = heapq.heappop(heap)
            if node in visited:
                continue
            visited.add(node)
            if len(visited) == n + 1:
                return current_total_time
            for nei, time in graph[node]:
                if nei not in visited:
                    heapq.heappush(heap, (current_total_time + time, nei))
        return -1

    def build_graph(self, edges):
        graph = defaultdict(list)
        for src, dest, t in edges:
            graph[src].append((dest, t))
            graph[dest].append((src, t))
        return graph

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

Đầu vào

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

Đầu ra

9