Computer >> Hướng Dẫn Máy Tính >  >> Lập Trình >> Python

Connecting Cities With Minimum Cost: Minimum Spanning Tree (Kruskal) Trong Python

Đề bài: Cho n thành phố và danh sách kết nối [city1, city2, cost]. Tìm chi phí nhỏ nhất để nối tất cả thành phố (đồ thị liên thông). Không thể thì trả -1.

Đây chính là bài toán Minimum Spanning Tree (MST) → thuật toán Kruskal + Union-Find (DSU).

Thuật toán Kruskal

  1. Sắp xếp các cạnh theo cost tăng dần.
  2. Duyệt cạnh: nếu 2 đỉnh thuộc component khác nhau → union, cộng cost.
  3. Dừng khi đã có n-1 cạnh (hoặc duyệt hết).
  4. Nếu component count > 1 → trả -1.

Code Python (Union-Find với Path Compression)


class Solution:
    def minimumCost(self, n: int, connections: list[list[int]]) -> int:
        parent = list(range(n + 1))
        rank = [0] * (n + 1)
        
        def find(x):
            if parent[x] != x:
                parent[x] = find(parent[x])
            return parent[x]
        
        def union(x, y):
            rx, ry = find(x), find(y)
            if rx == ry:
                return False
            if rank[rx] < rank[ry]:
                parent[rx] = ry
            elif rank[rx] > rank[ry]:
                parent[ry] = rx
            else:
                parent[ry] = rx
                rank[rx] += 1
            return True
        
        # Sort edges by cost
        connections.sort(key=lambda x: x[2])
        
        cost = 0
        edges_used = 0
        for u, v, w in connections:
            if union(u, v):
                cost += w
                edges_used += 1
                if edges_used == n - 1:
                    break
        
        return cost if edges_used == n - 1 else -1

# Test
ob = Solution()
print(ob.minimumCost(3, [[1,2,5], [1,3,6], [2,3,1]]))  # -1 (test case gốc có lỗi logic)
print(ob.minimumCost(3, [[1,2,5], [1,3,6], [2,3,1]]))  # Expected: 6 (chọn 1-3:6? wait)
# Sửa test: n=3, edges: 1-2:5, 1-3:6, 2-3:1 → MST chọn 2-3(1) + 1-2(5) = 6
print(ob.minimumCost(4, [[1,2,3],[3,4,4]]))  # -1 (không liên thông)

Lưu ý về test case gốc

Input: 3 [[1,2,5],[1,3,6],[2,3,1]] → MST cost = 1 + 5 = 6. Output gốc ghi -1 là sai (do code gốc có bug: disjoint = n-1 nhưng union logic sai).

Phân tích độ phức tạp

  • Time: O(E log E) do sort (E = số cạnh)
  • Space: O(n) cho parent/rank

Kết luận

Kruskal + DSU là giải pháp chuẩn cho MST. Union-Find với path compression + union by rank gần như O(1) amortized.