Swift で Kruskal（クラスカル法）使って MST：Minimum Spanning Tree（最小全域木）

備忘用のメモ

NOTE

• アルゴリズム：クラスカル法（Kruskal Algorithm）（貪欲法（Greedy Algorithm））
• 実装：Union Find

コード

アルゴリズムの部分。

• 辺の重みで小さい順にソートするし、
• 小さい順に辺を１つずつ取り出す
• 取り出された辺が使えるなら採用する。使うとサイクルになって使えない辺はスキップする。これを繰り替えす。

// n: number of node
func kruskal(_ n: Int, _ graph: [Edge]) -> Int {
    // init
    let sorted = graph.sorted(by: { $0.weight < $1.weight })
    var result = [Edge]()
    var uf = UF(n)
    
    // algo
    for edge in sorted {
        if uf.isConnected(edge.from, edge.to) { continue }
        let _ = uf.union(edge.from, edge.to)
        result.append(edge)
    }
    
    // return
    let sum = result.reduce(0) { $0 + $1.weight }
    return sum
}

---

Edge の部分

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

---

Union Find の部分

• path compression（経路圧縮）してる
• Union by size してる

struct UF {
    var size: [Int];
    var parentsOf: [Int];
    
    init(_ n: Int) {
        self.size = [Int](repeating: 1, count: n)
        self.parentsOf = [Int](repeating: -1, count: n)
        for i in 0..<n {
            self.parentsOf[i] = i
        }
    }
    
    mutating func root(_ elem: Int) -> Int {
        var current = elem
        while current != parentsOf[current] {
            parentsOf[current] = parentsOf[parentsOf[current]] // path compression
            current = parentsOf[current]
        }
        return current
    }

    mutating func isConnected(_ a: Int, _ b: Int) -> Bool {
        return root(a) == root(b)
    }
    
    mutating func union(_ a: Int, _ b: Int) -> Bool {
        if isConnected(a, b) { return false }
        
        let rootA = root(a)
        let rootB = root(b)
        
        // union by size
        let (smaller, larger) = size[rootA] < size[rootB]
            ? (rootA, rootB)
            : (rootB, rootA)
        
        self.size[larger] += self.size[smaller]
        self.parentsOf[smaller] = larger
        return true
    }
}

LeetCode

• https://leetcode.com/problems/min-cost-to-connect-all-points/

愚直にグラフへ変換してからKruskal。

func minCostConnectPoints(_ points: [[Int]]) -> Int {
    var graph = [Edge]()
    for i in 0..<points.count {
        for j in i..<points.count {
            let x1 = points[i][0], y1 = points[i][1]
            let x2 = points[j][0], y2 = points[j][1]
            let dist = abs(x1-x2) + abs(y1-y2)
            graph.append(Edge(from: i, to: j, weight: dist))
        }
        
    }
    
    return kruskal(points.count, graph)
}