[Heap] Change value of minimum or maximum element

Sources/PriorityQueueModule/Heap.swift

@@ -28,6 +28,23 @@ public struct Heap<Element: Comparable> {
   @usableFromInline
   internal var _storage: ContiguousArray<Element>
 
+  /// Creates a heap from the given sequence, assuming said sequence is already
+  /// an implicit data structure for a binary min-max heap.
+  ///
+  /// - Precondition: `storage` is finite, and it already matches the implicit
+  ///   data structure for a binary min-max heap.
+  ///
+  /// - Parameter storage: The elements of the heap.
+  /// - Postcondition: `unordered.elementsEqual(s)`, where *s* is a sequence
+  ///   with the same elements as pre-call `storage`.
+  ///
+  /// - Complexity: O(*n*), where *n* is the length of `storage`.
+  @_alwaysEmitIntoClient
+  internal init<S: Sequence>(_raw storage: S) where S.Element == Element {
+    _storage = .init(storage)
+    _checkInvariants()
+  }
+
   /// Creates an empty heap.
   @inlinable
   public init() {
@@ -169,6 +186,73 @@ extension Heap {
   public mutating func removeMax() -> Element {
     return popMax()!
   }
+
+  /// Changes the element with the lowest priority to the given value, then
+  /// redetermines which element in the heap has the lowest priority.
+  ///
+  /// - Precondition: `!isEmpty`.
+  ///
+  /// - Parameter replacement: The value overwriting the current lowest-priority
+  ///   element.
+  /// - Returns: The outgoing value of the overwritten element.
+  /// - Postcondition: When not equal, the multiplicity of the outgoing value
+  ///   will decrease by one, while the multiplicity of the incoming value will
+  ///   increase by one.  If the altered element now has a priority that is
+  ///   greater than any of the untargeted elements (if such exist), then
+  ///   `min()` will point to a different element.
+  ///
+  /// - Complexity: O(log `count`)
+  @inlinable @discardableResult
+  public mutating func replaceMin(with replacement: Element) -> Element {
+    precondition(!isEmpty, "No element to replace")
+
+    var removed = replacement
+    _update { handle in
+      let minNode = _Node.root
+      handle.swapAt(minNode, with: &removed)
+      handle.trickleDownMin(minNode)
+    }
+    _checkInvariants()
+    return removed
+  }
+
+  /// Changes the element with the highest priority to the given value, then
+  /// redetermines which element in the heap has the highest priority.
+  ///
+  /// - Precondition: `!isEmpty`.
+  ///
+  /// - Parameter replacement: The value overwriting the current
+  ///   highest-priority element.
+  /// - Returns: The outgoing value of the overwritten element.
+  /// - Postcondition: When not equal, the multiplicity of the outgoing value
+  ///   will decrease by one, while the multiplicity of the incoming value will
+  ///   increase by one.  If the altered element now has a priority that is less
+  ///   than any of the untargeted elements (if such exist), then `max()` will
+  ///   point to a different element.
+  ///
+  /// - Complexity: O(log `count`)
+  @inlinable @discardableResult
+  public mutating func replaceMax(with replacement: Element) -> Element {
+    precondition(!isEmpty, "No element to replace")
+
+    var removed = replacement
+    _update { handle in
+      switch handle.count {
+      case 1:
+        handle.swapAt(.root, with: &removed)
+      case 2:
+        handle.swapAt(.leftMax, with: &removed)
+        handle.bubbleUp(.leftMax)
+      default:
+        let maxNode = handle.maxValue(.leftMax, .rightMax)
+        handle.swapAt(maxNode, with: &removed)
+        handle.bubbleUp(maxNode)  // This must happen first
+        handle.trickleDownMax(maxNode)  // Either new element or dethroned min
+      }
+    }
+    _checkInvariants()
+    return removed
+  }
 }
 
 // MARK: -

Tests/PriorityQueueTests/HeapTests.swift

@@ -273,6 +273,93 @@ final class HeapTests: XCTestCase {
     XCTAssertEqual(heap.removeMax(), 1)
   }
 
+  func test_minimumReplacement() {
+    var heap = Heap(stride(from: 0, through: 27, by: 3).shuffled())
+    XCTAssertEqual(Array(heap.ascending), [0, 3, 6, 9, 12, 15, 18, 21, 24, 27])
+    XCTAssertEqual(heap.min(), 0)
+
+    // No change
+    heap.replaceMin(with: 0)
+    XCTAssertEqual(Array(heap.ascending), [0, 3, 6, 9, 12, 15, 18, 21, 24, 27])
+    XCTAssertEqual(heap.min(), 0)
+
+    // Even smaller
+    heap.replaceMin(with: -1)
+    XCTAssertEqual(Array(heap.ascending), [-1, 3, 6, 9, 12, 15, 18, 21, 24, 27])
+    XCTAssertEqual(heap.min(), -1)
+
+    // Larger, but not enough to usurp
+    heap.replaceMin(with: 2)
+    XCTAssertEqual(Array(heap.ascending), [2, 3, 6, 9, 12, 15, 18, 21, 24, 27])
+    XCTAssertEqual(heap.min(), 2)
+
+    // Larger, moving another element to be the smallest
+    heap.replaceMin(with: 5)
+    XCTAssertEqual(Array(heap.ascending), [3, 5, 6, 9, 12, 15, 18, 21, 24, 27])
+    XCTAssertEqual(heap.min(), 3)
+  }
+
+  func test_maximumReplacement() {
+    var heap = Heap(stride(from: 0, through: 27, by: 3).shuffled())
+    XCTAssertEqual(Array(heap.ascending), [0, 3, 6, 9, 12, 15, 18, 21, 24, 27])
+    XCTAssertEqual(heap.max(), 27)
+
+    // No change
+    heap.replaceMax(with: 27)
+    XCTAssertEqual(Array(heap.ascending), [0, 3, 6, 9, 12, 15, 18, 21, 24, 27])
+    XCTAssertEqual(heap.max(), 27)
+
+    // Even larger
+    heap.replaceMax(with: 28)
+    XCTAssertEqual(Array(heap.ascending), [0, 3, 6, 9, 12, 15, 18, 21, 24, 28])
+    XCTAssertEqual(heap.max(), 28)
+
+    // Smaller, but not enough to usurp
+    heap.replaceMax(with: 26)
+    XCTAssertEqual(Array(heap.ascending), [0, 3, 6, 9, 12, 15, 18, 21, 24, 26])
+    XCTAssertEqual(heap.max(), 26)
+
+    // Smaller, moving another element to be the largest
+    heap.replaceMax(with: 23)
+    XCTAssertEqual(Array(heap.ascending), [0, 3, 6, 9, 12, 15, 18, 21, 23, 24])
+    XCTAssertEqual(heap.max(), 24)
+
+    // Check the finer details.  As these peek into the stored structure, they
+    // may need to be updated whenever the internal format changes.
+    var heap2 = Heap(_raw: [1])
+    XCTAssertEqual(heap2.max(), 1)
+    XCTAssertEqual(Array(heap2.unordered), [1])
+    XCTAssertEqual(heap2.replaceMax(with: 2), 1)
+    XCTAssertEqual(heap2.max(), 2)
+    XCTAssertEqual(Array(heap2.unordered), [2])
+
+    heap2 = Heap(_raw: [1, 2])
+    XCTAssertEqual(heap2.max(), 2)
+    XCTAssertEqual(Array(heap2.unordered), [1, 2])
+    XCTAssertEqual(heap2.replaceMax(with: 3), 2)
+    XCTAssertEqual(heap2.max(), 3)
+    XCTAssertEqual(Array(heap2.unordered), [1, 3])
+    XCTAssertEqual(heap2.replaceMax(with: 0), 3)
+    XCTAssertEqual(heap2.max(), 1)
+    XCTAssertEqual(Array(heap2.unordered), [0, 1])
+
+    heap2 = Heap(_raw: [5, 20, 31, 16, 8, 7, 18])
+    XCTAssertEqual(heap2.max(), 31)
+    XCTAssertEqual(Array(heap2.unordered), [5, 20, 31, 16, 8, 7, 18])
+    XCTAssertEqual(heap2.replaceMax(with: 29), 31)
+    XCTAssertEqual(Array(heap2.unordered), [5, 20, 29, 16, 8, 7, 18])
+    XCTAssertEqual(heap2.max(), 29)
+    XCTAssertEqual(heap2.replaceMax(with: 19), 29)
+    XCTAssertEqual(Array(heap2.unordered), [5, 20, 19, 16, 8, 7, 18])
+    XCTAssertEqual(heap2.max(), 20)
+    XCTAssertEqual(heap2.replaceMax(with: 15), 20)
+    XCTAssertEqual(Array(heap2.unordered), [5, 16, 19, 15, 8, 7, 18])
+    XCTAssertEqual(heap2.max(), 19)
+    XCTAssertEqual(heap2.replaceMax(with: 4), 19)
+    XCTAssertEqual(Array(heap2.unordered), [4, 16, 18, 15, 8, 7, 5])
+    XCTAssertEqual(heap2.max(), 18)
+  }
+
   // MARK: -
 
   func test_min_struct() {