# 测试代码
from mst import Graph

if __name__ == "__main__":
    g = Graph(5)  # 创建一个包含5个顶点的图
    g.add_edge(0, 1, 2)
    g.add_edge(0, 3, 6)
    g.add_edge(1, 3, 8)
    g.add_edge(1, 2, 3)
    g.add_edge(1, 4, 5)
    g.add_edge(2, 4, 7)
    g.add_edge(3, 4, 9)

    print("Prim 算法结果:")
    cost, edges = g.prim_mst()
    print(f"总权值: {cost}, 边: {edges}")

    print("\nKruskal 算法结果:")
    cost, edges = g.kruskal_mst()
    print(f"总权值: {cost}, 边: {edges}")