四十三庵

蔀の雑記帳

Kruskal's Algorithm

実装

func minimumSpanningTreeKruskal<T>(graph: Graph<T>) -> (cost: Int, tree: Graph<T>) {
    var cost: Int = 0
    var tree = Graph<T>()
    let sortedEdgeListByWeight = graph.edgeList.sorted(by: { $0.weight < $1.weight })
    
    var unionFind = UnionFind<T>()
    for vertex in graph.vertices {
        unionFind.addSetWith(vertex)
    }
    
    for edge in sortedEdgeListByWeight {
        let v1 = edge.vertex1
        let v2 = edge.vertex2
        if !unionFind.inSameSet(v1, and: v2) {
            cost += edge.weight
            tree.addEdge(edge)
            unionFind.unionSetsContaining(v1, and: v2)
        }
    }
    
    return (cost: cost, tree: tree)
}

ポイント

  • UnionFindの使い方で詰まった
    • 吐かせるとこんな感じ UnionFind(index: [6: 3, 4: 4, 3: 5, 5: 2, 1: 0, 2: 1], parent: [0, 1, 2, 3, 4, 5], size: [1, 1, 1, 1, 1, 1]
    • graph.verticesはsortされていないので、ランダムに入る
    • indexは辞書型になっていて、たとえばVertex1番は、parent[0]と紐づいていて、自分自身を親に持つ
    • これが集合になると、parent: [0, 0, 2, 3, 4, 5], size: [2, 1, 1, 1, 1, 1]みたいになる
  • コスト最小の辺の両端のVertexをunionFindで集合にする
  • そうするとunionFind.inSameSet(v1, and: v2)で既にMinumum Spanning Tree内で連結されたvertexは上手いこと除外できる
  • UnionFindの使い方がわかれば、そんなに難しくないか