Prims.swift

//===--- Prims.swift ------------------------------------------------------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//

// The test implements Prim's algorithm for minimum spanning tree building.
// http://en.wikipedia.org/wiki/Prim%27s_algorithm

// This class implements array-based heap (priority queue).
// It is used to store edges from nodes in spanning tree to nodes outside of it.
// We are interested only in the edges with the smallest costs, so if there are
// several edges pointing to the same node, we keep only one from them. Thus,
// it is enough to record this node instead.
// We maintain a map (node index in graph)->(node index in heap) to be able to
// update the heap fast when we add a new node to the tree.
import TestsUtils

class PriorityQueue {
  final var heap: Array<EdgeCost>
  final var graphIndexToHeapIndexMap: Array<Int?>

  // Create heap for graph with NUM nodes.
  init(Num: Int) {
    heap = Array<EdgeCost>()
    graphIndexToHeapIndexMap = Array<Int?>(repeating:nil, count: Num)
  }

  func isEmpty() -> Bool {
    return heap.isEmpty
  }

  // Insert element N to heap, maintaining the heap property.
  func insert(_ n: EdgeCost) {
    let ind: Int = heap.count
    heap.append(n)
    graphIndexToHeapIndexMap[n.to] = heap.count - 1
    bubbleUp(ind)
  }

  // Insert element N if in's not in the heap, or update its cost if the new
  // value is less than the existing one.
  func insertOrUpdate(_ n: EdgeCost) {
    let id = n.to
    let c  = n.cost
    if let ind = graphIndexToHeapIndexMap[id] {
      if heap[ind].cost <= c {
        // We don't need an edge with a bigger cost
        return
      }
      heap[ind].cost = c
      heap[ind].from = n.from
      bubbleUp(ind)
    } else {
      insert(n)
    }
  }

  // Restore heap property by moving element at index IND up.
  // This is needed after insertion, and after decreasing an element's cost.
  func bubbleUp(_ ind: Int) {
    var ind = ind
    let c = heap[ind].cost
    while (ind != 0) {
      let p = getParentIndex(ind)
      if heap[p].cost > c {
        Swap(p, with: ind)
        ind = p
      } else {
        break
      }
    }
  }

  // Pop minimum element from heap and restore the heap property after that.
  func pop() -> EdgeCost? {
    if (heap.isEmpty) {
      return nil
    }
    Swap(0, with:heap.count-1)
    let r = heap.removeLast()
    graphIndexToHeapIndexMap[r.to] = nil
    bubbleDown(0)
    return r
  }

  // Restore heap property by moving element at index IND down.
  // This is needed after removing an element, and after increasing an
  // element's cost.
  func bubbleDown(_ ind: Int) {
    var ind = ind
    let n = heap.count
    while (ind < n) {
      let l = getLeftChildIndex(ind)
      let r = getRightChildIndex(ind)
      if (l >= n) {
        break
      }
      var min: Int
      if (r < n && heap[r].cost < heap[l].cost) {
        min = r
      } else {
        min = l
      }
      if (heap[ind].cost <= heap[min].cost) {
        break
      }
      Swap(ind, with: min)
      ind = min
    }
  }

  // Swaps elements I and J in the heap and correspondingly updates
  // graphIndexToHeapIndexMap.
  func Swap(_ i: Int, with j : Int) {
    if (i == j) {
      return
    }
    (heap[i], heap[j]) = (heap[j], heap[i])
    let (I, J) = (heap[i].to, heap[j].to)
    (graphIndexToHeapIndexMap[I], graphIndexToHeapIndexMap[J]) =
    (graphIndexToHeapIndexMap[J], graphIndexToHeapIndexMap[I])
  }

  // Dumps the heap.
  func dump() {
    print("QUEUE")
    for nodeCost in heap {
      let to: Int = nodeCost.to
      let from: Int = nodeCost.from
      let cost: Double = nodeCost.cost
      print("(\(from)->\(to), \(cost))")
    }
  }

  func getLeftChildIndex(_ index : Int) -> Int {
    return index*2 + 1
  }
  func getRightChildIndex(_ index : Int) -> Int {
    return (index + 1)*2
  }
  func getParentIndex(_ childIndex : Int) -> Int {
    return (childIndex - 1)/2
  }
}

struct GraphNode {
  var id: Int
  var adjList: Array<Int>

  init(i : Int) {
    id = i
    adjList = Array<Int>()
  }
}

struct EdgeCost {
  var to: Int
  var cost: Double
  var from: Int
}

struct Edge : Equatable {
  var start: Int
  var end: Int
}

func ==(lhs: Edge, rhs: Edge) -> Bool {
  return lhs.start == rhs.start && lhs.end == rhs.end
}

extension Edge : Hashable {
  func hash(into hasher: inout Hasher) {
    hasher.combine(start)
    hasher.combine(end)
  }
}

func Prims(_ graph : Array<GraphNode>, _ fun : (Int, Int) -> Double) -> Array<Int?> {
  var treeEdges = Array<Int?>(repeating:nil, count:graph.count)

  let queue = PriorityQueue(Num:graph.count)
  // Make the minimum spanning tree root its own parent for simplicity.
  queue.insert(EdgeCost(to: 0, cost: 0.0, from: 0))

  // Take an element with the smallest cost from the queue and add its
  // neighbors to the queue if their cost was updated
  while !queue.isEmpty() {
    // Add an edge with minimum cost to the spanning tree
    let e = queue.pop()!
    let newnode = e.to
    // Add record about the edge newnode->e.from to treeEdges
    treeEdges[newnode] = e.from

    // Check all adjacent nodes and add edges, ending outside the tree, to the
    // queue. If the queue already contains an edge to an adjacent node, we
    // replace existing one with the new one in case the new one costs less.
    for adjNodeIndex in graph[newnode].adjList {
      if treeEdges[adjNodeIndex] != nil {
        continue
      }
      let newcost = fun(newnode, graph[adjNodeIndex].id)
      queue.insertOrUpdate(EdgeCost(to: adjNodeIndex, cost: newcost, from: newnode))
    }
  }
  return treeEdges
}