Minimum Spanning Tree (Kruskal's and Prim's)

Author: Xiangxin

Minimum Spanning Tree/Images/kruskal.png

@@ -0,0 +1,31 @@
+//
+//  Kruskal.swift
+//  
+//
+//  Created by xiang xin on 16/3/17.
+//
+//
+
+
+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)
+}

Minimum Spanning Tree/Images/prim.png

@@ -0,0 +1,93 @@
+/**
+ The following code demonstrates getting MST of the graph below by both
+ Kruskal's and Prim's algorithms.
+ */
+
+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)
+}
+
+func minimumSpanningTreePrim<T>(graph: Graph<T>) -> (cost: Int, tree: Graph<T>) {
+  var cost: Int = 0
+  var tree = Graph<T>()
+
+  if graph.vertices.isEmpty {
+    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: graph.vertices.first!, weight: 0, parent: nil))
+  while let head = priorityQueue.dequeue() {
+    let vertex = head.vertex
+    if visited.contains(vertex) {
+      continue
+    }
+    visited.insert(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) {
+          priorityQueue.enqueue((vertex: nextVertex, weight: neighbour.weight, parent: vertex))
+        }
+      }
+    }
+  }
+
+  return (cost: cost, tree: tree)
+}
+
+/*:
+ ![Graph](mst.png)
+ */
+
+var graph = Graph<Int>()
+graph.addEdge(vertex1: 1, vertex2: 2, weight: 6)
+graph.addEdge(vertex1: 1, vertex2: 3, weight: 1)
+graph.addEdge(vertex1: 1, vertex2: 4, weight: 5)
+graph.addEdge(vertex1: 2, vertex2: 3, weight: 5)
+graph.addEdge(vertex1: 2, vertex2: 5, weight: 3)
+graph.addEdge(vertex1: 3, vertex2: 4, weight: 5)
+graph.addEdge(vertex1: 3, vertex2: 5, weight: 6)
+graph.addEdge(vertex1: 3, vertex2: 6, weight: 4)
+graph.addEdge(vertex1: 4, vertex2: 6, weight: 2)
+graph.addEdge(vertex1: 5, vertex2: 6, weight: 6)
+
+print("===== Kruskal's =====")
+let result1 = minimumSpanningTreeKruskal(graph: graph)
+print("Minimum spanning tree total weight: \(result1.cost)")
+print("Minimum spanning tree:")
+print(result1.tree)
+
+print("===== Prim's =====")
+let result2 = minimumSpanningTreePrim(graph: graph)
+print("Minimum spanning tree total weight: \(result2.cost)")
+print("Minimum spanning tree:")
+print(result2.tree)

Minimum Spanning Tree/Kruskal.swift

@@ -0,0 +1,55 @@
+
+// Undirected edge
+public struct Edge<T>: CustomStringConvertible {
+  public let vertex1: T
+  public let vertex2: T
+  public let weight: Int
+
+  public var description: String {
+    return "[\(vertex1)-\(vertex2), \(weight)]"
+  }
+}
+
+
+// Undirected weighted graph
+public struct Graph<T: Hashable>: CustomStringConvertible {
+
+  public private(set) var edgeList: [Edge<T>]
+  public private(set) var vertices: Set<T>
+  public private(set) var adjList: [T: [(vertex: T, weight: Int)]]
+
+  public init() {
+    edgeList = [Edge<T>]()
+    vertices = Set<T>()
+    adjList = [T: [(vertex: T, weight: Int)]]()
+  }
+
+  public var description: String {
+    var description = ""
+    for edge in edgeList {
+      description += edge.description + "\n"
+    }
+    return description
+  }
+
+  public mutating func addEdge(vertex1 v1: T, vertex2 v2: T, weight w: Int) {
+    edgeList.append(Edge(vertex1: v1, vertex2: v2, weight: w))
+    vertices.insert(v1)
+    vertices.insert(v2)
+    if adjList[v1] == nil {
+      adjList[v1] = [(vertex: v2, weight: w)]
+    } else {
+      adjList[v1]!.append((vertex: v2, weight: w))
+    }
+
+    if adjList[v2] == nil {
+      adjList[v2] = [(vertex: v1, weight: w)]
+    } else {
+      adjList[v2]!.append((vertex: v1, weight: w))
+    }
+  }
+
+  public mutating func addEdge(_ edge: Edge<T>) {
+    addEdge(vertex1: edge.vertex1, vertex2: edge.vertex2, weight: edge.weight)
+  }
+}