Swift で Prim（プリム法）使って MST：Minimum Spanning Tree（最小全域木）

備忘用メモ

Note

• アルゴリズム：Prim's Algorithm
• 実装：Primary Queue（Heap）

コード

アルゴリズム部分。

• どれか１つ頂点から決めで始点を決める。始点から生えている辺をすべてエンキューする。
• デキューする。PriorityQueueなのでデキューされたものはキューイングされたものの中で一番重さの値が低い辺になる。
• デキューされた辺の相手側の頂点を主役に取り、ここから生えている辺をすべてキューイングする。これを繰り返す。

// n: number of node
func prim(_ n: Int, _ graph: [[Edge]]) -> Int {
    // init
    var heap = MinHeap.init()
    var mst = [Edge]()
    var visited = [Bool](repeating: false, count: n)

    // first node
    let from = 0
    for edge in graph[from] { heap.push(item: edge) }
    visited[from] = true

    // loop
    while let edge = heap.pop() {
        if visited[edge.to] { continue }
        visited[edge.to] = true
        mst.append(edge)
        for next in graph[edge.to] {
            heap.push(item: next)
        }
    }

    // return
    var sum = 0
    for edge in mst { sum += edge.weight }
    return sum
}

---

Primary Queue となる Min Heap の実装部分

struct MinHeap {
    // zero-based index
    var list = [Edge]()

    private mutating func heapifyUp(_ idx: Int) {
        if !self.hasParent(idx) { return }
        if parentOf(idx).weight > list[idx].weight {
            let parentIdx = parentIdxOf(idx)
            list.swapAt(parentIdx, idx)
            self.heapifyUp(parentIdx)
        }
    }
    private mutating func heapifyDown(_ idx: Int) {
        if !self.hasLeft(idx) { return }
        let smallerIdx: Int = {
            if !self.hasRight(idx) { return self.leftIdxOf(idx)}
            return self.leftOf(idx).weight < self.rightOf(idx).weight
                ? self.leftIdxOf(idx)
                : self.rightIdxOf(idx)
        }()
        if list[idx].weight < list[smallerIdx].weight { return }
        list.swapAt(idx, smallerIdx)
        heapifyDown(smallerIdx)
    }

    private func parentIdxOf(_ idx: Int) -> Int { Int(floor((Double(idx) - 1) / 2)) }
    private func leftIdxOf(_ idx: Int)  -> Int { idx*2 + 1 }
    private func rightIdxOf(_ idx: Int) -> Int { idx*2 + 2 }
    private func parentOf(_ idx: Int) -> Edge { list[self.parentIdxOf(idx)] }
    private func leftOf(_ idx: Int)  -> Edge { list[self.leftIdxOf(idx)] }
    private func rightOf(_ idx: Int) -> Edge { list[self.rightIdxOf(idx)] }
    private func hasParent(_ idx: Int) -> Bool { idx != 0 }
    private func hasLeft(_ idx: Int) -> Bool { self.leftIdxOf(idx) < list.count }
    private func hasRight(_ idx: Int) -> Bool { self.rightIdxOf(idx) < list.count }

    mutating func push(item: Edge) {
        list.append(item)
        heapifyUp(list.count-1)
    }

    mutating func pop() -> Edge? {
        if list.isEmpty { return nil }
        if list.count == 1 { return list.removeFirst() }
        let peaked = list[0]
        list[0] = list[list.count-1]
        list.removeLast()
        heapifyDown(0)
        return peaked
    }
}

---

Edge のデータ構造

struct Edge {
    let to: Int
    let weight: Int
}

REF

• https://leetcode.com/problems/min-cost-to-connect-all-points/
  • この記事のアルゴリズムだと Time Exceeded になるので注意