四十三庵

蔀の雑記帳

Prim’s Algorithm

  • 計算量はO(n2)
  • 解説はここがわかりやすかった

実装

func minimumSpanningTreePrim<T>(graph: Graph<T>) -> (cost: Int, tree: Graph<T>) {
    var cost: Int = 0
    var tree = Graph<T>()
    
    guard let start = graph.vertices.first else {
        return (cost: cost, tree: tree)
    }
    
    var visited = Set<T>()
    var priorityQueue = PriorityQueue<(vertex: T, weight: Int, parent: T?)>(
        sort: { $0.weight < $1.weight })
    
    priorityQueue.enqueue((vertex: start, weight: 0, parent: nil))
    while let head = priorityQueue.dequeue() {
        let vertex = head.vertex
        if visited.contains(vertex) {
            continue
        }
        visited.insert(vertex)
        print("visited: ", vertex)
        
        cost += head.weight
        if let prev = head.parent {
            tree.addEdge(vertex1: prev, vertex2: vertex, weight: head.weight)
        }
        
        if let neighbours = graph.adjList[vertex] {
            for neighbour in neighbours {
                let nextVertex = neighbour.vertex
                if !visited.contains(nextVertex) {
                    print("nextVertex", nextVertex)
                    print("weight", neighbour.weight)
                    priorityQueue.enqueue((vertex: nextVertex, weight: neighbour.weight, parent: vertex))
                }
            }
        }
    }
    
    return (cost: cost, tree: tree)
}

引っかかったポイント

  • graph.vertices.firstはランダム
  • priorityQueueは探索対象となったVertexに隣接するVertex間のEdgeのCostで優先度をつける
  • adjListで上手いこと隣接するVertex&Costを取れるので、それをVisitedでない限りpriorityQueueに突っ込む
  • すると次のenqueueする対象は最もコストの低い経路を選ぶ
  • 更にpriorityQueueには選択されなかった経路も残ったまま、次の探索対象の隣接Vertexを突っ込む
  • 全てのVertexが探索し終わると、集合visitedに全Vertexが入るので、そのまま無限にdequeueしてループを抜ける