- Adopted discussed improvements

- Made nodes and some other methods internal, so that extensions like HeapSort can use them. This fixes the HeapSort algorithm as well as the tests for both.

Heap Sort/HeapSort.swift

@@ -1,10 +1,10 @@
 extension Heap {
   public mutating func sort() -> [T] {
-    for i in stride(from: (elements.count - 1), through: 1, by: -1) {
-      elements.swapAt(0, i)
-      shiftDown(0, heapSize: i)
+    for i in stride(from: (nodes.count - 1), through: 1, by: -1) {
+      nodes.swapAt(0, i)
+      shiftDown(from: 0, until: i)
     }
-    return elements
+    return nodes
   }
 }
 

Heap/Heap.swift

@@ -6,7 +6,7 @@
 public struct Heap<T> {
 
   /** The array that stores the heap's nodes. */
-  private(set) var nodes = [T]()
+  internal var nodes = [T]()
 
   /**
    * Determines how to compare two nodes in the heap.
@@ -33,14 +33,14 @@ public struct Heap<T> {
    */
   public init(array: [T], sort: @escaping (T, T) -> Bool) {
     self.orderCriteria = sort
-    buildHeap(fromArray: array)
+    configureHeap(from: array)
   }
 
   /**
-   * Creates the max-heap or min-heap from an array, in a bottom-up manner.
+   * Configures the max-heap or min-heap from an array, in a bottom-up manner.
    * Performance: This runs pretty much in O(n).
    */
-  private mutating func buildHeap(fromArray array: [T]) {
+  private mutating func configureHeap(from array: [T]) {
     nodes = array
     for i in stride(from: (nodes.count/2-1), through: 0, by: -1) {
       shiftDown(i)
@@ -59,7 +59,7 @@ public struct Heap<T> {
    * Returns the index of the parent of the element at index i.
    * The element at index 0 is the root of the tree and has no parent.
    */
-  @inline(__always) func parentIndex(ofIndex i: Int) -> Int {
+  @inline(__always) internal func parentIndex(ofIndex i: Int) -> Int {
     return (i - 1) / 2
   }
 
@@ -68,7 +68,7 @@ public struct Heap<T> {
    * Note that this index can be greater than the heap size, in which case
    * there is no left child.
    */
-  @inline(__always) func leftChildIndex(ofIndex i: Int) -> Int {
+  @inline(__always) internal func leftChildIndex(ofIndex i: Int) -> Int {
     return 2*i + 1
   }
 
@@ -77,7 +77,7 @@ public struct Heap<T> {
    * Note that this index can be greater than the heap size, in which case
    * there is no right child.
    */
-  @inline(__always) func rightChildIndex(ofIndex i: Int) -> Int {
+  @inline(__always) internal func rightChildIndex(ofIndex i: Int) -> Int {
     return 2*i + 2
   }
 
@@ -89,12 +89,6 @@ public struct Heap<T> {
     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).
@@ -121,7 +115,7 @@ public struct Heap<T> {
   public mutating func replace(index i: Int, value: T) {
     guard i < nodes.count else { return }
 
-    removeAt(i)
+    remove(at: i)
     insert(value)
   }
 
@@ -130,26 +124,25 @@ public struct Heap<T> {
    * value; for a min-heap it is the minimum value. Performance: O(log n).
    */
   @discardableResult public mutating func remove() -> T? {
-    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
-      }
+    guard !nodes.isEmpty else { return nil }
+
+    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).
    * Note that you need to know the node's index.
    */
-  @discardableResult public mutating func removeAt(_ index: Int) -> T? {
+  @discardableResult public mutating func remove(at index: Int) -> T? {
     guard index < nodes.count else { return nil }
 
     let size = nodes.count - 1
@@ -165,7 +158,7 @@ public struct Heap<T> {
    * Takes a child node and looks at its parents; if a parent is not larger
    * (max-heap) or not smaller (min-heap) than the child, we exchange them.
    */
-  mutating func shiftUp(_ index: Int) {
+  internal mutating func shiftUp(_ index: Int) {
     var childIndex = index
     let child = nodes[childIndex]
     var parentIndex = self.parentIndex(ofIndex: childIndex)
@@ -183,7 +176,7 @@ public struct Heap<T> {
    * Looks at a parent node and makes sure it is still larger (max-heap) or
    * smaller (min-heap) than its childeren.
    */
-  private mutating func shiftDown(from index: Int, until endIndex: Int) {
+  internal mutating func shiftDown(from index: Int, until endIndex: Int) {
     let leftChildIndex = self.leftChildIndex(ofIndex: index)
     let rightChildIndex = leftChildIndex + 1
 
@@ -204,7 +197,7 @@ public struct Heap<T> {
     shiftDown(from: first, until: endIndex)
   }
 
-  private mutating func shiftDown(_ index: Int) {
+  internal mutating func shiftDown(_ index: Int) {
     shiftDown(from: index, until: nodes.count)
   }
 
@@ -222,7 +215,7 @@ extension Heap where T: Equatable {
   /** 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 remove(at: index)
     }
     return nil
   }

Heap/Tests/HeapTests.swift

@@ -154,14 +154,14 @@ class HeapTests: XCTestCase {
     XCTAssertEqual(h4.peek()!, 0)
   }
 
-  func testCreateMaxHeapEqualElements() {
+  func testCreateMaxHeapEqualnodes() {
     let heap = Heap(array: [1, 1, 1, 1, 1], sort: >)
     XCTAssertTrue(verifyMaxHeap(heap))
     XCTAssertTrue(verifyMinHeap(heap))
     XCTAssertEqual(heap.nodes, [1, 1, 1, 1, 1])
   }
 
-  func testCreateMinHeapEqualElements() {
+  func testCreateMinHeapEqualnodes() {
     let heap = Heap(array: [1, 1, 1, 1, 1], sort: <)
     XCTAssertTrue(verifyMinHeap(heap))
     XCTAssertTrue(verifyMaxHeap(heap))
@@ -198,33 +198,33 @@ class HeapTests: XCTestCase {
     }
   }
 
-  func testRemovingAtIndex() {
+  func testRemoving() {
     var h = Heap(array: [100, 50, 70, 10, 20, 60, 65], sort: >)
     XCTAssertTrue(verifyMaxHeap(h))
     XCTAssertEqual(h.nodes, [100, 50, 70, 10, 20, 60, 65])
 
     //test index out of bounds
-    let v = h.removeAt(10)
+    let v = h.remove(at: 10)
     XCTAssertEqual(v, nil)
     XCTAssertTrue(verifyMaxHeap(h))
     XCTAssertEqual(h.nodes, [100, 50, 70, 10, 20, 60, 65])
 
-    let v1 = h.removeAt(5)
+    let v1 = h.remove(at: 5)
     XCTAssertEqual(v1, 60)
     XCTAssertTrue(verifyMaxHeap(h))
     XCTAssertEqual(h.nodes, [100, 50, 70, 10, 20, 65])
 
-    let v2 = h.removeAt(4)
+    let v2 = h.remove(at: 4)
     XCTAssertEqual(v2, 20)
     XCTAssertTrue(verifyMaxHeap(h))
     XCTAssertEqual(h.nodes, [100, 65, 70, 10, 50])
 
-    let v3 = h.removeAt(4)
+    let v3 = h.remove(at: 4)
     XCTAssertEqual(v3, 50)
     XCTAssertTrue(verifyMaxHeap(h))
     XCTAssertEqual(h.nodes, [100, 65, 70, 10])
 
-    let v4 = h.removeAt(0)
+    let v4 = h.remove(at: 0)
     XCTAssertEqual(v4, 100)
     XCTAssertTrue(verifyMaxHeap(h))
     XCTAssertEqual(h.nodes, [70, 65, 10])
@@ -267,17 +267,6 @@ class HeapTests: XCTestCase {
     XCTAssertEqual(h.nodes, [13, 12, 9, 5, 6, 8, 7, 4, 0, 1, 2])
   }
 
-  func testRemoveNode() {
-    var h = Heap(array: [15, 13, 9, 5, 12, 8, 7, 4, 0, 6, 2, 1], sort: >)
-    XCTAssertTrue(verifyMaxHeap(h))
-    XCTAssertEqual(h.nodes, [15, 13, 9, 5, 12, 8, 7, 4, 0, 6, 2, 1])
-    XCTAssertEqual(h.node(at: 3)!, 5)
-    let v = h.remove(node: 5)
-    XCTAssertEqual(v, 5)
-    XCTAssertTrue(verifyMaxHeap(h))
-    XCTAssertFalse(h.nodes.contains(5))
-  }
-
   func testRemoveRandomItems() {
     for n in 1...40 {
       var a = randomArray(n)
@@ -288,16 +277,8 @@ class HeapTests: XCTestCase {
       let m = (n + 1)/2
       for k in 1...m {
         let i = Int(arc4random_uniform(UInt32(n - k + 1)))
-
-        var v: Int? = nil
-        if k == 2 || k == m {
-          v = h.remove(node: h.node(at: i)!)
-        } else {
-          v = h.removeAt(i)
-        }
-        XCTAssertNotNil(v)
-
-        let j = a.index(of: v!)!
+        let v = h.remove(at: i)!
+        let j = a.index(of: v)!
         a.remove(at: j)
 
         XCTAssertTrue(verifyMaxHeap(h))