Merge pull request kodecocodes#685 from DXXL/master

Improvements to Heap

Author: kelvinlauKL

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

@@ -4,65 +4,62 @@
 //
 
 public struct Heap<T> {
+  
   /** The array that stores the heap's nodes. */
-  var elements = [T]()
+  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
-    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)
-   }
-   }
-   */
+    self.orderCriteria = sort
+    configureHeap(from: array)
+  }
 
   /**
-   * Converts an array to a max-heap or min-heap 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).
    */
-  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 configureHeap(from 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
   }
 
   /**
    * 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
   }
 
@@ -71,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
   }
 
@@ -80,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,139 +86,138 @@ public struct Heap<T> {
    * value (for a min-heap).
    */
   public func peek() -> T? {
-    return elements.first
+    return nodes.first
   }
 
   /**
    * 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)
+    remove(at: 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()
+    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 = elements[0]
-      elements[0] = elements.removeLast()
-      shiftDown()
+      let value = nodes[0]
+      nodes[0] = nodes.removeLast()
+      shiftDown(0)
       return value
     }
   }
 
   /**
-   * 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 remove(at 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()
   }
 
   /**
    * 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 = 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
+  internal 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)
   }
+  
+  internal 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 remove(at: 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)
     }
   }