Improvements to Heap

- renamed "elements" to "nodes" to clarify the difference between array and heap. This also unifies code and comments, which were sometimes referring to the array as "nodes".
- made nodes publicly readable but only settable within the heap itself. To change the nodes/operate on the heap, the provided functions should be used, to guarantee the heaps structure.
- renamed isOrderedBefore to orderCriteria. The variable name isOrderedBefore did not clearly explain what the variable is for/what the closure stored within does
- added method to get node at a specific index
- replace method should not force the replaced value to be higher/lower than its parent. Instead replace should remove the node to be replaced and insert the given value as a new one in order to preserve the heap's structure. Asserting crashes the whole program.
- added @discardableResult to removeAt. If remove is @discardableResult removeAt should be as well in for consistency.
- replaced while true in shiftDown with a recursive call. while true is bad programming style.
- simplified indexOf method
- added a removeNode method for heaps with equatable data type
- improved, adjusted and added descriptions
- made functions private where fileprivate was not needed
- added tests for new functions and tested everything
- updated README where necessary (to conform to new code, only a few changes)

Author: davisriedel

Heap/Heap.swift

@@ -4,58 +4,55 @@
 //
 
 public struct Heap<T> {
+  
   /** The array that stores the heap's nodes. */
-  var elements = [T]()
+  private(set) var nodes = [T]()
 
-  /** Determines whether this is a max-heap (>) or min-heap (<). */
-  fileprivate var isOrderedBefore: (T, T) -> Bool
+  /**
+   * Determines how to compare two nodes in the heap.
+   * Use '>' for a max-heap or '<' for a min-heap,
+   * or provide a comparing method if the heap is made
+   * of custom elements, for example tuples.
+   */
+  private var orderCriteria: (T, T) -> Bool
 
   /**
    * Creates an empty heap.
    * The sort function determines whether this is a min-heap or max-heap.
-   * For integers, > makes a max-heap, < makes a min-heap.
+   * For comparable data types, > makes a max-heap, < makes a min-heap.
    */
   public init(sort: @escaping (T, T) -> Bool) {
-    self.isOrderedBefore = sort
+    self.orderCriteria = sort
   }
 
   /**
    * Creates a heap from an array. The order of the array does not matter;
    * the elements are inserted into the heap in the order determined by the
-   * sort function.
+   * sort function. For comparable data types, '>' makes a max-heap,
+   * '<' makes a min-heap.
    */
   public init(array: [T], sort: @escaping (T, T) -> Bool) {
-    self.isOrderedBefore = sort
+    self.orderCriteria = sort
     buildHeap(fromArray: array)
   }
 
-  /*
-   // This version has O(n log n) performance.
-   private mutating func buildHeap(array: [T]) {
-   elements.reserveCapacity(array.count)
-   for value in array {
-   insert(value)
-   }
-   }
-   */
-  
   /**
-   * Converts an array to a max-heap or min-heap in a bottom-up manner.
+   * Creates the max-heap or min-heap from an array, in a bottom-up manner.
    * Performance: This runs pretty much in O(n).
    */
-  fileprivate mutating func buildHeap(fromArray array: [T]) {
-    elements = array
-    for i in stride(from: (elements.count/2 - 1), through: 0, by: -1) {
-      shiftDown(i, heapSize: elements.count)
+  private mutating func buildHeap(fromArray array: [T]) {
+    nodes = array
+    for i in stride(from: (nodes.count/2-1), through: 0, by: -1) {
+      shiftDown(i)
     }
   }
 
   public var isEmpty: Bool {
-    return elements.isEmpty
+    return nodes.isEmpty
   }
 
   public var count: Int {
-    return elements.count
+    return nodes.count
   }
 
   /**
@@ -89,69 +86,79 @@ public struct Heap<T> {
    * value (for a min-heap).
    */
   public func peek() -> T? {
-    return elements.first
+    return nodes.first
+  }
+  
+  /** Returns the node at given index */
+  public func node(at i: Int) -> T? {
+    guard i < nodes.count else { return nil }
+    return nodes[i]
   }
 
   /**
    * Adds a new value to the heap. This reorders the heap so that the max-heap
    * or min-heap property still holds. Performance: O(log n).
    */
   public mutating func insert(_ value: T) {
-    elements.append(value)
-    shiftUp(elements.count - 1)
+    nodes.append(value)
+    shiftUp(nodes.count - 1)
   }
 
+  /**
+   * Adds a sequence of values to the heap. This reorders the heap so that
+   * the max-heap or min-heap property still holds. Performance: O(log n).
+   */
   public mutating func insert<S: Sequence>(_ sequence: S) where S.Iterator.Element == T {
     for value in sequence {
       insert(value)
     }
   }
 
   /**
-   * Allows you to change an element. In a max-heap, the new element should be
-   * larger than the old one; in a min-heap it should be smaller.
+   * Allows you to change an element. This reorders the heap so that
+   * the max-heap or min-heap property still holds.
    */
   public mutating func replace(index i: Int, value: T) {
-    guard i < elements.count else { return }
+    guard i < nodes.count else { return }
 
-    assert(isOrderedBefore(value, elements[i]))
-    elements[i] = value
-    shiftUp(i)
+    removeAt(i)
+    insert(value)
   }
 
   /**
    * Removes the root node from the heap. For a max-heap, this is the maximum
    * value; for a min-heap it is the minimum value. Performance: O(log n).
    */
   @discardableResult public mutating func remove() -> T? {
-    if elements.isEmpty {
-      return nil
-    } else if elements.count == 1 {
-      return elements.removeLast()
-    } else {
-      // Use the last node to replace the first one, then fix the heap by
-      // shifting this new first node into its proper position.
-      let value = elements[0]
-      elements[0] = elements.removeLast()
-      shiftDown()
-      return value
+    if !nodes.isEmpty {
+      if nodes.count == 1 {
+        return nodes.removeLast()
+      } else {
+        // Use the last node to replace the first one, then fix the heap by
+        // shifting this new first node into its proper position.
+        let value = nodes[0]
+        nodes[0] = nodes.removeLast()
+        shiftDown(0)
+        return value
+      }
     }
+    return nil
   }
 
   /**
-   * Removes an arbitrary node from the heap. Performance: O(log n). You need
-   * to know the node's index, which may actually take O(n) steps to find.
+   * Removes an arbitrary node from the heap. Performance: O(log n).
+   * Note that you need to know the node's index.
    */
-  public mutating func removeAt(_ index: Int) -> T? {
-    guard index < elements.count else { return nil }
+  @discardableResult public mutating func removeAt(_ index: Int) -> T? {
+    guard index < nodes.count else { return nil }
 
-    let size = elements.count - 1
+    let size = nodes.count - 1
     if index != size {
-      elements.swapAt(index, size)
-      shiftDown(index, heapSize: size)
+      nodes.swapAt(index, size)
+      shiftDown(from: index, until: size)
       shiftUp(index)
     }
-    return elements.removeLast()
+    return nodes.removeLast()
   }
 
   /**
@@ -160,68 +167,64 @@ public struct Heap<T> {
    */
   mutating func shiftUp(_ index: Int) {
     var childIndex = index
-    let child = elements[childIndex]
+    let child = nodes[childIndex]
     var parentIndex = self.parentIndex(ofIndex: childIndex)
 
-    while childIndex > 0 && isOrderedBefore(child, elements[parentIndex]) {
-      elements[childIndex] = elements[parentIndex]
+    while childIndex > 0 && orderCriteria(child, nodes[parentIndex]) {
+      nodes[childIndex] = nodes[parentIndex]
       childIndex = parentIndex
       parentIndex = self.parentIndex(ofIndex: childIndex)
     }
 
-    elements[childIndex] = child
-  }
-  
-  mutating func shiftDown() {
-    shiftDown(0, heapSize: elements.count)
+    nodes[childIndex] = child
   }
 
   /**
    * Looks at a parent node and makes sure it is still larger (max-heap) or
    * smaller (min-heap) than its childeren.
    */
-  mutating func shiftDown(_ index: Int, heapSize: Int) {
-    var parentIndex = index
+  private mutating func shiftDown(from index: Int, until endIndex: Int) {
+    let leftChildIndex = self.leftChildIndex(ofIndex: index)
+    let rightChildIndex = leftChildIndex + 1
 
-    while true {
-      let leftChildIndex = self.leftChildIndex(ofIndex: parentIndex)
-      let rightChildIndex = leftChildIndex + 1
-      
-      // Figure out which comes first if we order them by the sort function:
-      // the parent, the left child, or the right child. If the parent comes
-      // first, we're done. If not, that element is out-of-place and we make
-      // it "float down" the tree until the heap property is restored.
-      var first = parentIndex
-      if leftChildIndex < heapSize && isOrderedBefore(elements[leftChildIndex], elements[first]) {
-        first = leftChildIndex
-      }
-      if rightChildIndex < heapSize && isOrderedBefore(elements[rightChildIndex], elements[first]) {
-        first = rightChildIndex
-      }
-      if first == parentIndex { return }
-      
-      elements.swapAt(parentIndex, first)
-      parentIndex = first
+    // Figure out which comes first if we order them by the sort function:
+    // the parent, the left child, or the right child. If the parent comes
+    // first, we're done. If not, that element is out-of-place and we make
+    // it "float down" the tree until the heap property is restored.
+    var first = index
+    if leftChildIndex < endIndex && orderCriteria(nodes[leftChildIndex], nodes[first]) {
+      first = leftChildIndex
+    }
+    if rightChildIndex < endIndex && orderCriteria(nodes[rightChildIndex], nodes[first]) {
+      first = rightChildIndex
     }
+    if first == index { return }
+    
+    nodes.swapAt(index, first)
+    shiftDown(from: first, until: endIndex)
+  }
+  
+  private mutating func shiftDown(_ index: Int) {
+    shiftDown(from: index, until: nodes.count)
   }
+  
 }
 
 // MARK: - Searching
 
 extension Heap where T: Equatable {
-  /**
-   * Searches the heap for the given element. Performance: O(n).
-   */
-  public func index(of element: T) -> Int? {
-    return index(of: element, 0)
+  
+  /** Get the index of a node in the heap. Performance: O(n). */
+  public func index(of node: T) -> Int? {
+    return nodes.index(where: { $0 == node })
   }
 
-  fileprivate func index(of element: T, _ i: Int) -> Int? {
-    if i >= count { return nil }
-    if isOrderedBefore(element, elements[i]) { return nil }
-    if element == elements[i] { return i }
-    if let j = index(of: element, self.leftChildIndex(ofIndex: i)) { return j }
-    if let j = index(of: element, self.rightChildIndex(ofIndex: i)) { return j }
+  /** Removes the first occurrence of a node from the heap. Performance: O(n log n). */
+  @discardableResult public mutating func remove(node: T) -> T? {
+    if let index = index(of: node) {
+      return removeAt(index)
+    }
     return nil
   }
+  
 }

Heap/README.markdown

@@ -38,7 +38,7 @@ A heap is not a replacement for a binary search tree, and there are similarities
 
 **Balancing.** A binary search tree must be "balanced" so that most operations have **O(log n)** performance. You can either insert and delete your data in a random order or use something like an [AVL tree](../AVL%20Tree/) or [red-black tree](../Red-Black%20Tree/), but with heaps we don't actually need the entire tree to be sorted. We just want the heap property to be fulfilled, so balancing isn't an issue. Because of the way the heap is structured, heaps can guarantee **O(log n)** performance.
 
-**Searching.** Whereas searching is fast in a binary tree, it is slow in a heap. Searching isn't a top priority in a heap since the purpose of a heap is to put the largest (or smallest) node at the front and to allow relatively fast inserts and deletes. 
+**Searching.** Whereas searching is fast in a binary tree, it is slow in a heap. Searching isn't a top priority in a heap since the purpose of a heap is to put the largest (or smallest) node at the front and to allow relatively fast inserts and deletes.
 
 ## The tree inside an array
 
@@ -148,7 +148,7 @@ There are two primitive operations necessary to make sure the heap is a valid ma
 
 Shifting up or down is a recursive procedure that takes **O(log n)** time.
 
-Here are other operations that are built on primitive operations: 
+Here are other operations that are built on primitive operations:
 
 - `insert(value)`: Adds the new element to the end of the heap and then uses `shiftUp()` to fix the heap.
 
@@ -190,7 +190,7 @@ The `(16)` was added to the first available space on the last row.
 
 Unfortunately, the heap property is no longer satisfied because `(2)` is above `(16)`, and we want higher numbers above lower numbers. (This is a max-heap.)
 
-To restore the heap property, we swap `(16)` and `(2)`. 
+To restore the heap property, we swap `(16)` and `(2)`.
 
 ![The heap before insertion](Images/Insert2.png)
 
@@ -214,7 +214,7 @@ What happens to the empty spot at the top?
 
 ![The root is gone](Images/Remove1.png)
 
-When inserting, we put the new value at the end of the array. Here, we do the opposite: we take the last object we have, stick it up on top of the tree, and restore the heap property. 
+When inserting, we put the new value at the end of the array. Here, we do the opposite: we take the last object we have, stick it up on top of the tree, and restore the heap property.
 
 ![The last node goes to the root](Images/Remove2.png)
 
@@ -226,7 +226,7 @@ Keep shifting down until the node does not have any children or it is larger tha
 
 ![The last node goes to the root](Images/Remove4.png)
 
-The time required for shifting all the way down is proportional to the height of the tree which takes **O(log n)** time. 
+The time required for shifting all the way down is proportional to the height of the tree which takes **O(log n)** time.
 
 > **Note:** `shiftUp()` and `shiftDown()` can only fix one out-of-place element at a time. If there are multiple elements in the wrong place, you need to call these functions once for each of those elements.
 
@@ -279,7 +279,7 @@ In code:
 ```swift
   private mutating func buildHeap(fromArray array: [T]) {
     elements = array
-    for i in (elements.count/2 - 1).stride(through: 0, by: -1) {
+    for i in stride(from: (nodes.count/2-1), through: 0, by: -1) {
       shiftDown(index: i, heapSize: elements.count)
     }
   }