Chương trình tìm cây khung nhỏ nhất (MST) bằng thuật toán Prim trong Python

Giả sử cho một đồ thị, cần tìm Cây Khung Nhỏ Nhất (MST). MST là tập con của đồ thị có trọng số chứa tất cả đỉnh, liên thông, không chu trình, và tổng trọng số cạnh nhỏ nhất. Ở đây ta dùng thuật toán Prim để tìm tổng trọng số cạnh của MST.

Ví dụ, với đầu vào như hình minh họa, số đỉnh n = 4, đỉnh bắt đầu s = 3, đầu ra sẽ là 14.

MST từ đồ thị này sẽ như hình minh họa, tổng trọng số cạnh là 14.

Giải pháp

Để giải quyết bài toán này, ta thực hiện các bước sau:

• Định nghĩa hàm mst_find(G, s):
  • distance := danh sách kích thước G khởi tạo infinity
  • distance[s] := 0
  • itr := danh sách kích thước G khởi tạo False
  • c := 0
  • Lặp vô hạn:
    • min_weight := infinity, m_idx := -1
    • Với i từ 0 đến kích thước G:
      • Nếu itr[i] == False và distance[i] < min_weight:
        • min_weight := distance[i], m_idx := i
    • Nếu m_idx == -1: break
    • c := c + min_weight
    • itr[m_idx] := True
    • Với mỗi cặp i, j trong G[m_idx].items(): distance[i] := min(distance[i], j)
  • Trả về c
• G := map mới chứa n map khác
• Với mỗi item trong edges:
  • u := item[0]-1, v := item[1]-1, w := item[2]
  • min_weight = min(G[u][v], w) nếu tồn tại, ngược lại w
  • G[u][v] := min_weight, G[v][u] := min_weight
• Trả về mst_find(G, s)

Ví dụ minh họa

Xem triển khai sau để hiểu rõ hơn:

def mst_find(G, s):
    distance = [float("inf")] * len(G)
    distance[s] = 0
    itr = [False] * len(G)
    c = 0
    while True:
        min_weight = float("inf")
        m_idx = -1
        for i in range(len(G)):
            if itr[i] == False:
                if distance[i] < min_weight:
                    min_weight = distance[i]
                    m_idx = i
        if m_idx == -1:
            break
        c += min_weight
        itr[m_idx] = True
        for i, j in G[m_idx].items():
            distance[i] = min(distance[i], j)
    return c

def solve(n, edges, s):
    G = {i: {} for i in range(n)}
    for item in edges:
        u = item[0] - 1
        v = item[1] - 1
        w = item[2]
        try:
            min_weight = min(G[u][v], w)
            G[u][v] = min_weight
            G[v][u] = min_weight
        except KeyError:
            G[u][v] = w
            G[v][u] = w
    return mst_find(G, s)

print(solve(4, [(1, 2, 5), (1, 3, 5), (2, 3, 7), (1, 4, 4)], 3))

Đầu vào

4, [(1, 2, 5), (1, 3, 5), (2, 3, 7), (1, 4, 4)], 3

Đầu ra

14