Skip to content

Instantly share code, notes, and snippets.

@CTMacUser
Last active July 18, 2020 04:28
Show Gist options
  • Save CTMacUser/c6fe95f0b6b6eda9f558ead4a29d01cb to your computer and use it in GitHub Desktop.
Save CTMacUser/c6fe95f0b6b6eda9f558ead4a29d01cb to your computer and use it in GitHub Desktop.
Heap operation methods and Priority Queues in Swift, inspired by the C++ Standard Library. See <https://forums.swift.org/t/heaps-structure-not-storage-and-priority-queues/31135> for more information.
extension Collection {
/// Computes the binary-tree children indices for the given starting index,
/// if the collection is long enough.
///
/// Child elements of a flattened tree are consecutive, so if two indices
/// are returned, `distance(from: left!, to: right!) == 1`.
///
/// - Parameter source: The index of the parent element.
/// - Returns: A tuple with the indices of the left and right binary-tree
/// child elements for the element at `source`. If the collection is
/// short enough to only contain one, `right` will be `nil`. If the
/// collection is too short to support either, both `left` and `right`
/// will be `nil`.
///
/// - Complexity: O(1) for random-access collections; O(*n*) otherwise,
/// where *n* is the length of the collection.
@usableFromInline
func binaryHeapChildrenIndices(of source: Index) -> (left: Index?, right: Index?) {
let start = startIndex, end = endIndex
let sourceOffset = distance(from: start, to: source)
guard
let leftChild = index(start, offsetBy: 2 * sourceOffset + 1, limitedBy: end),
leftChild < end else {
return (nil, nil)
}
let rightChild = index(after: leftChild)
return (leftChild, rightChild < end ? rightChild : nil)
}
/// Computes the binary-tree children index range of the given starting
/// index, if the collection is long enough.
///
/// - Parameter source: The index of the parent element.
/// - Returns: The index range for both the left and right binary-tree
/// child elements for the element at `source`. If the collection is
/// short enough for only the left child, a one-element suffix range is
/// returned. If the collection is too short for either child, an empty
/// range at `endIndex` is returned.
///
/// - Complexity: O(1) for random-access collections; O(*n*) otherwise,
/// where *n* is the length of the collection.
@usableFromInline
func binaryHeapChildrenRange(of source: Index) -> Range<Index> {
let (possibleLeft, possibleRight) = binaryHeapChildrenIndices(of: source)
let end = endIndex
guard let leftChild = possibleLeft else { return end..<end }
guard let rightChild = possibleRight else { return leftChild..<end }
assert(distance(from: leftChild, to: rightChild) == 1)
return leftChild..<index(after: rightChild)
}
/// Computes the binary-tree parent index of the given starting index.
///
/// The root index is its own parent index.
///
/// - Precondition: `source < endIndex`.
///
/// - Parameter source: The index of a child element.
/// - Returns: The index of the earlier element that would have `source` as
/// either its left or right binary-tree child index.
///
/// - Complexity: O(1) for random-access collections; O(*n*) otherwise,
/// where *n* is the length of the collection.
@usableFromInline
func heapParentIndex(of source: Index) -> Index {
let start = startIndex
let sourceOffset = distance(from: start, to: source)
return index(start, offsetBy: (sourceOffset - 1) / 2)
}
}
extension RandomAccessCollection {
/// Computes where the collection stops being a heap of the given priority,
/// using the given predicate to compare between elements.
///
/// The predicate must be a strict weak ordering over the elements.
///
/// - Parameter priority: Whether the lowest- or highest-ordered element is
/// the lead of the heap.
/// - Parameter areInIncreasingOrder: A predicate that returns `true` if its
/// first argument should be ordered before its second argument;
/// otherwise, `false`.
/// - Returns: The index of the first element that is disordered relative to
/// its heap-graph parent, or `endIndex` if no elements meet that
/// criterion.
///
/// - Complexity: O(*n*), where *n* is the length of the collection.
@inlinable
public func endIndexOfHeap(prioritizing priority: HeapPriority, by areInIncreasingOrder: (Element, Element) throws -> Bool) rethrows -> Index {
for parent in indices {
let (possibleLeft, possibleRight) = binaryHeapChildrenIndices(of: parent)
assert(possibleLeft != nil || possibleRight == nil)
guard let left = possibleLeft else { break }
guard try priority.areProperlyHeaped(parent: self[parent], child: self[left], by: areInIncreasingOrder) else { return left }
guard let right = possibleRight else { break }
guard try priority.areProperlyHeaped(parent: self[parent], child: self[right], by: areInIncreasingOrder) else { return right }
assert(left < right)
}
return endIndex
}
/// Determines if the collection is a heap of the given priority, using the
/// given predicate to compare between elements.
///
/// The predicate must be a strict weak ordering over the elements.
///
/// - Parameter priority: Whether the lowest- or highest-ordered element is
/// the lead of the heap.
/// - Parameter areInIncreasingOrder: A predicate that returns `true` if its
/// first argument should be ordered before its second argument;
/// otherwise, `false`.
/// - Returns: Whether the entire collection conforms to a heap.
///
/// - Complexity: O(*n*), where *n* is the length of the collection.
@inlinable
public func isHeap(prioritizing priority: HeapPriority, by areInIncreasingOrder: (Element, Element) throws -> Bool) rethrows -> Bool {
return try endIndexOfHeap(prioritizing: priority, by: areInIncreasingOrder) == endIndex
}
}
extension RandomAccessCollection where Element: Comparable {
/// Computes where the collection stops being a binary min heap.
///
/// - Returns: The index of the first element that is less than its
/// heap-graph parent, or `endIndex` if no such element exists.
///
/// - Complexity: O(*n*), where *n* is the length of the collection.
@inlinable
public func endIndexOfMinHeap() -> Index {
return endIndexOfHeap(prioritizing: .min, by: <)
}
/// Computes where the collection stops being a binary max heap.
///
/// - Returns: The index of the first element that is greater than its
/// heap-graph parent, or `endIndex` if no such element exists.
///
/// - Complexity: O(*n*), where *n* is the length of the collection.
@inlinable
public func endIndexOfMaxHeap() -> Index {
return endIndexOfHeap(prioritizing: .max, by: <)
}
/// Determines if the collection is a binary min heap.
///
/// - Returns: Whether the entire collection conforms to a min heap.
///
/// - Complexity: O(*n*), where *n* is the length of the collection.
@inlinable
public func isMinHeap() -> Bool {
return isHeap(prioritizing: .min, by: <)
}
/// Determines if the collection is a binary max heap.
///
/// - Returns: Whether the entire collection conforms to a max heap.
///
/// - Complexity: O(*n*), where *n* is the length of the collection.
@inlinable
public func isMaxHeap() -> Bool {
return isHeap(prioritizing: .max, by: <)
}
}
extension MutableCollection where Self: RandomAccessCollection {
/// Assimilates the element following a prefix subsequence heap of the given
/// priority into the heap, using the given predicate to compare between
/// elements.
///
/// The predicate must be a strict weak ordering over the elements.
///
/// - Precondition: `self[..<pivot].isHeap(priority, areInIncreasingOrder)`,
/// `pivot < endIndex`.
///
/// - Parameter pivot: The location of the element to insert into the heap,
/// and the past-the-end index for that heap.
/// - Parameter priority: Whether the lowest- or highest-ordered element is
/// the lead of the heap.
/// - Parameter areInIncreasingOrder: A predicate that returns `true` if its
/// first argument should be ordered before its second argument;
/// otherwise, `false`.
/// - Returns: The index after `pivot`, signifying the end of the expanded
/// heap.
///
/// - Postcondition: `self[...pivot].isHeap(priority, areInIncreasingOrder)`.
///
/// - Complexity: O(log *n*), where *n* is the length of the collection.
@discardableResult
public mutating func partitionIntoHeap(at pivot: Index, prioritizing priority: HeapPriority, by areInIncreasingOrder: (Element, Element) throws -> Bool) rethrows -> Index {
let start = startIndex
var current = pivot
repeat {
let parent = heapParentIndex(of: current)
guard try !priority.areProperlyHeaped(parent: self[parent], child: self[current], by: areInIncreasingOrder) else { break }
swapAt(current, parent)
current = parent
} while current > start
return index(after: pivot)
}
/// Checks if the sub-heap starting at the given root index really conforms
/// to a heap of the given priority, using the given predicate to compare
/// elements, then swaps elements around the sub-heap until it conforms if
/// necessary.
///
/// The predicate must be a strict weak ordering over the elements.
///
/// - Precondition: `subHeapRoot` must be dereferenceable. If so, and it
/// has children (*i.e.* sub-sub-heaps), those children must already
/// conform as compatible heaps.
///
/// - Parameter subHeapRoot: The top of the sub-heap to search over.
/// - Parameter priority: Whether the lowest- or highest-ordered element is
/// the lead of the heap.
/// - Parameter areInIncreasingOrder: A predicate that returns `true` if its
/// first argument should be ordered before its second argument;
/// otherwise, `false`.
/// - Returns: The final location for the element that was at `subHeapRoot`
/// when this method started.
///
/// - Postcondition: The sub-heap starting at `subHeapRoot` conforms as a
/// compatible heap.
///
/// - Complexity: O(log *n*), where *n* is the length of the collection.
@usableFromInline
@discardableResult
mutating func confirmHeap(at subHeapRoot: Index, prioritizing priority: HeapPriority, by areInIncreasingOrder: (Element, Element) throws -> Bool) rethrows -> Index {
precondition(subHeapRoot < endIndex)
var root = subHeapRoot, children = [Index]()
let rootValue = self[root]
children.reserveCapacity(2)
while true {
children.replaceSubrange(children.startIndex..., with: indices[binaryHeapChildrenRange(of: root)])
guard try !children.allSatisfy({ try priority.areProperlyHeaped(parent: rootValue, child: self[$0], by: areInIncreasingOrder) }) else {
return root
}
let nextRoot: Index!
assert(!children.isEmpty)
switch priority {
case .min:
nextRoot = try children.min { try areInIncreasingOrder(self[$0], self[$1]) }
case .max:
nextRoot = try children.max { try areInIncreasingOrder(self[$0], self[$1]) }
}
swapAt(root, nextRoot)
root = nextRoot
}
}
/// Moves the lead element of this heap to the end and reorders the
/// remaining elements as a prefix that is still a heap of the given
/// priority, using the given predicate to compare elements.
///
/// The predicate must be a strict weak ordering over the elements.
///
/// - Precondition: `isHeap(priority, areInIncreasingOrder)`.
///
/// - Parameter priority: Whether the lowest- or highest-ordered element is
/// the lead of the heap.
/// - Parameter areInIncreasingOrder: A predicate that returns `true` if its
/// first argument should be ordered before its second argument;
/// otherwise, `false`.
/// - Returns: The index of the former heap-leading element, which is also
/// the end index of the shrunken heap.
///
/// - Postcondition: `dropLast().isHeap(priority, areInIncreasingOrder)`.
///
/// - Complexity: O(log *n*), where *n* is the length of the collection.
@inlinable
@discardableResult
public mutating func partitionOutHeapLead(prioritizing priority: HeapPriority, by areInIncreasingOrder: (Element, Element) throws -> Bool) rethrows -> Index {
let start = startIndex
guard let lastIndex = index(endIndex, offsetBy: -1, limitedBy: start), lastIndex > start else { return start }
swapAt(start, lastIndex)
try self[..<lastIndex].confirmHeap(at: start, prioritizing: priority, by: areInIncreasingOrder)
return lastIndex
}
/// Reorders the elements of the collection into a heap with the given
/// priority, where either the least- or most-ordered element, using the
/// given predicate to compare elements, is sent to the lead.
///
/// The predicate must be a strict weak ordering over the elements.
///
/// - Parameter priority: Whether the lowest- or highest-ordered element is
/// the lead of the heap.
/// - Parameter areInIncreasingOrder: A predicate that returns `true` if its
/// first argument should be ordered before its second argument;
/// otherwise, `false`.
///
/// - Postcondition: `isHeap(priority, areInIncreasingOrder)`.
///
/// - Complexity: O(*n*), where *n* is the length of the collection.
@inlinable
public mutating func arrangeIntoHeap(prioritizing priority: HeapPriority, by areInIncreasingOrder: (Element, Element) throws -> Bool) rethrows {
// Collections with at most one element are already heaped.
guard count >= 2 else { return }
let lastIndexWithChildren = index(startIndex, offsetBy: count / 2 - 1)
try uponEachIndexBackwards { innerSelf, index in
guard index <= lastIndexWithChildren else { return }
try innerSelf.confirmHeap(at: index, prioritizing: priority, by: areInIncreasingOrder)
}
}
/// If this heap of the given priority is not empty, then update its lead
/// element with the given action closure and move said element if necessary
/// to maintain the heap property, using the given predicate to compare
/// elements.
///
/// The predicate must be a strict weak ordering over the elements.
///
/// - Precondition: `isHeap(priority, areInIncreasingOrder)`.
///
/// - Parameter body: A closure to mutate its given element.
/// - Parameter element: The (lead) element to be mutated.
/// - Parameter priority: Whether the lowest- or highest-ordered element is
/// the lead of the heap.
/// - Parameter areInIncreasingOrder: A predicate that returns `true` if its
/// first argument should be ordered before its second argument;
/// otherwise, `false`.
/// - Returns: `nil` if the lead element coming into this method stays as
/// the lead; otherwise, the new location of the former lead element.
///
/// - Postcondition: `isHeap(priority, areInIncreasingOrder)`, but
/// corresponding values aren't necessarily the same from before the call.
///
/// - Complexity: O(log *n*), where *n* is the length of the collection.
@inlinable
@discardableResult
public mutating func updateHeapLead(using body: (_ element: inout Element) throws -> Void, prioritizing priority: HeapPriority, by areInIncreasingOrder: (Element, Element) throws -> Bool) rethrows -> Index? {
let start = startIndex
guard start < endIndex else { return nil }
try body(&self[start])
let newIndexForLead = try confirmHeap(at: start, prioritizing: priority, by: areInIncreasingOrder)
return newIndexForLead == start ? nil : newIndexForLead
}
/// Sorts the elements of the collection from a heap with the given
/// priority, using the given predicate to compare elements.
///
/// The predicate must be a strict weak ordering over the elements.
///
/// - Parameter priority: Whether the lowest- or highest-ordered element is
/// the lead of the heap.
/// - Parameter areInIncreasingOrder: A predicate that returns `true` if its
/// first argument should be ordered before its second argument;
/// otherwise, `false`.
///
/// - Postcondition: When `priority == .max`,
/// `isSorted(areInIncreasingOrder)`; otherwise,
/// `reversed().isSorted(areInIncreasingOrder)`.
///
/// - Complexity: O(*n* × log *n*), where *n* is the length of the collection.
@inlinable
public mutating func sortOutHeap(prioritizing priority: HeapPriority, by areInIncreasingOrder: (Element, Element) throws -> Bool) rethrows {
var heapEnd = endIndex
let secondIndex = index(startIndex, offsetBy: +1, limitedBy: heapEnd) ?? heapEnd
while heapEnd > secondIndex {
// The returned end goes down by one each turn.
heapEnd = try self[..<heapEnd].partitionOutHeapLead(prioritizing: priority, by: areInIncreasingOrder)
}
}
}
extension MutableCollection where Self: RandomAccessCollection, Element: Comparable {
/// Assimilates the element following a prefix subsequence min heap into
/// said heap.
///
/// - Precondition: `self[..<pivot].isMinHeap()`, `pivot < endIndex`.
///
/// - Parameter pivot: The location of the element to insert into the heap,
/// and the past-the-end index for that heap.
/// - Returns: The index after `pivot`, signifying the end of the expanded
/// heap.
///
/// - Postcondition: `self[...pivot].isMinHeap()`.
///
/// - Complexity: O(log *n*), where *n* is the length of the collection.
@inlinable
@discardableResult
public mutating func partitionIntoMinHeap(at pivot: Index) -> Index {
return partitionIntoHeap(at: pivot, prioritizing: .min, by: <)
}
/// Moves the lead element of this min heap to the end and reorders the
/// remaining elements as a prefix that is still a heap.
///
/// - Precondition: `isMinHeap()`.
///
/// - Returns: The index of the former heap-leading element, which is also
/// the end index of the shrunken heap.
///
/// - Postcondition: `dropLast().isMinHeap()`.
///
/// - Complexity: O(log *n*), where *n* is the length of the collection.
@inlinable
@discardableResult
public mutating func partitionOutMinHeapLead() -> Index {
return partitionOutHeapLead(prioritizing: .min, by: <)
}
/// Reorders the elements of the collection into a min heap, where the
/// least-ordered element is sent to the lead.
///
/// - Postcondition: `isMinHeap()`.
///
/// - Complexity: O(*n*), where *n* is the length of the collection.
@inlinable
public mutating func arrangeIntoMinHeap() {
arrangeIntoHeap(prioritizing: .min, by: <)
}
/// If this min heap is not empty, then update its lead element with the
/// given action closure and move said element if necessary to maintain the
/// heap property.
///
/// - Precondition: `isMinHeap()`.
///
/// - Parameter body: A closure to mutate its given element.
/// - Parameter element: The (lead) element to be mutated.
/// - Returns: `nil` if the lead element coming into this method stays as
/// the lead; otherwise, the new location of the former lead element.
///
/// - Postcondition: `isMinHeap()`, but corresponding values aren't
/// necessarily the same from before the call.
///
/// - Complexity: O(log *n*), where *n* is the length of the collection.
@inlinable
@discardableResult
public mutating func updateMinHeapLead(using body: (_ element: inout Element) throws -> Void) rethrows -> Index? {
return try updateHeapLead(using: body, prioritizing: .min, by: <)
}
/// Sorts the elements of the collection from a min heap to a monotonically
/// decreasing sequence.
///
/// - Postcondition: `reversed().isSorted()`.
///
/// - Complexity: O(*n* × log *n*), where *n* is the length of the collection.
@inlinable
public mutating func sortOutMinHeap() {
sortOutHeap(prioritizing: .min, by: <)
}
/// Assimilates the element following a prefix subsequence max heap into
/// said heap.
///
/// - Precondition: `self[..<pivot].isMaxHeap()`, `pivot < endIndex`.
///
/// - Parameter pivot: The location of the element to insert into the heap,
/// and the past-the-end index for that heap.
/// - Returns: The index after `pivot`, signifying the end of the expanded
/// heap.
///
/// - Postcondition: `self[...pivot].isMaxHeap()`.
///
/// - Complexity: O(log *n*), where *n* is the length of the collection.
@inlinable
@discardableResult
public mutating func partitionIntoMaxHeap(at pivot: Index) -> Index {
return partitionIntoHeap(at: pivot, prioritizing: .max, by: <)
}
/// Moves the lead element of this max heap to the end and reorders the
/// remaining elements as a prefix that is still a heap.
///
/// - Precondition: `isMaxHeap()`.
///
/// - Returns: The index of the former heap-leading element, which is also
/// the end index of the shrunken heap.
///
/// - Postcondition: `dropLast().isMaxHeap()`.
///
/// - Complexity: O(log *n*), where *n* is the length of the collection.
@inlinable
@discardableResult
public mutating func partitionOutMaxHeapLead() -> Index {
return partitionOutHeapLead(prioritizing: .max, by: <)
}
/// Reorders the elements of the collection into a max heap, where the
/// most-ordered element is sent to the lead.
///
/// - Postcondition: `isMaxHeap()`.
///
/// - Complexity: O(*n*), where *n* is the length of the collection.
@inlinable
public mutating func arrangeIntoMaxHeap() {
arrangeIntoHeap(prioritizing: .max, by: <)
}
/// If this max heap is not empty, then update its lead element with the
/// given action closure and move said element if necessary to maintain the
/// heap property.
///
/// - Precondition: `isMaxHeap()`.
///
/// - Parameter body: A closure to mutate its given element.
/// - Parameter element: The (lead) element to be mutated.
/// - Returns: `nil` if the lead element coming into this method stays as
/// the lead; otherwise, the new location of the former lead element.
///
/// - Postcondition: `isMaxHeap()`, but corresponding values aren't
/// necessarily the same from before the call.
///
/// - Complexity: O(log *n*), where *n* is the length of the collection.
@inlinable
@discardableResult
public mutating func updateMaxHeapLead(using body: (_ element: inout Element) throws -> Void) rethrows -> Index? {
return try updateHeapLead(using: body, prioritizing: .max, by: <)
}
/// Sorts the elements of the collection from a max heap to a monotonically
/// increasing sequence.
///
/// - Postcondition: `isSorted()`.
///
/// - Complexity: O(*n* × log *n*), where *n* is the length of the collection.
@inlinable
public mutating func sortOutMaxHeap() {
sortOutHeap(prioritizing: .max, by: <)
}
}
/// Indicator of which extreme value a heap should prioritize.
public enum HeapPriority {
/// The lowest-ordered element is set as the heap's lead.
case min
/// The highest-ordered element is set as the heap's lead.
case max
}
extension HeapPriority {
/// Determines if the values of the given parent and child elements respect
/// what is needed to maintain a heap of the receiver's priority, using the
/// given predicate to compare elements.
///
/// The predicate must be a strict weak ordering over the elements.
///
/// - Parameter parent: The element that will be the lead of the heap.
/// - Parameter child: The element that should trail within the heap.
/// - Parameter areInIncreasingOrder: A predicate that returns `true` if its
/// first argument should be ordered before its second argument;
/// otherwise, `false`.
/// - Returns: Whether the value of `parent` is in the right direction for
/// the heap relative to the value of `child`. (If the two elements are
/// equivalent, `true` is returned.)
@inlinable
public func areProperlyHeaped<T>(parent: T, child: T, by areInIncreasingOrder: (T, T) throws -> Bool) rethrows -> Bool {
switch self {
case .min:
// The parent has to rank lower (or equal).
return try !areInIncreasingOrder(child, parent)
case .max:
// The parent has to rank higher (or equal).
return try !areInIncreasingOrder(parent, child)
}
}
}
/// A collection adapter controlling access to only the best-ranked element.
public struct SomePriorityQueue<Base: MutableCollection & RandomAccessCollection> {
/// The type of vended elements.
public typealias Element = Base.Element
/// The container for the elements.
@usableFromInline
var base: Base
/// The predicate for ranking element values.
@usableFromInline
let areInIncreasingOrder: (Element, Element) -> Bool
/// Whether to publish the element with the lowest or highest rank.
public let priority: HeapPriority
/// Creates a priority queue from the given collection, with its ordering
/// checked by the given predicate, vending in the order of the given
/// priority.
///
/// - Precondition: `areInIncreasingOrder` must be a strict weak ordering
/// over the elements.
///
/// - Parameter base: The source of the initial elements.
/// - Parameter priority: Indicates which end of the ordering scale that
/// elements will be read from.
/// - Parameter areInIncreasingOrder: A predicate that returns `true` if its
/// first argument should be ordered before its second argument;
/// otherwise, `false`.
///
/// - Postcondition: The values of `base` are rearranged to ease the vending
/// of either the lowest ranked (when `priority == .min`) or highest
/// ranked (when `priority == .max`) elements as determined by
/// `areInIncreasingOrder`.
///
/// - Complexity: O(*n*), where *n* is the length of `base`.
@inlinable
public init(base: Base, priority: HeapPriority, by areInIncreasingOrder: @escaping (Element, Element) -> Bool) {
self.base = base
self.priority = priority
self.areInIncreasingOrder = areInIncreasingOrder
self.base.arrangeIntoHeap(prioritizing: priority, by: areInIncreasingOrder)
}
}
extension SomePriorityQueue {
/// Reveals the best-ranked element, if any.
///
/// The best-ranked element is the minimal value when `priority` is `.min`,
/// and the maximal value when `priority` is `.max`.
@inlinable
public var first: Element? { base.first }
/// The number of elements in the queue.
@inlinable
public var count: Int { base.count }
/// Whether the queue is empty.
@inlinable
public var isEmpty: Bool { base.isEmpty }
/// Returns the elements of the queue, sorted.
///
/// - Warning: Do not use this method if `Base` is a reference type.
///
/// - Returns: A copy of the elements of this collection. When `priority`
/// is `.max`, the sort order is monotonically increasing as determined by
/// the ordering predicate. Otherwise, the sort is monotonically
/// decreasing.
///
/// - Complexity: O(*n* × log *n*), where *n* is the length of the queue.
@inlinable
public func sorted() -> Base {
var copy = base
copy.sortOutHeap(prioritizing: priority, by: areInIncreasingOrder)
return copy
}
/// Mutates the currently best-ranked element with the given closure, moving
/// it if its updated value disqualifies it from staying the best.
///
/// - Parameter body: A closure that can mutate its passed-in element.
/// - Returns: `true` if a different element is promoted to be the
/// best-ranked element; otherwise `false` if the incoming best-ranked
/// element stays the best, or if the queue is empty.
///
/// - Postcondition: If the queue is not empty, the best-ranked element
/// changes, and at least one hidden element out-ranks the updated value,
/// then the queue's elements are rearranged.
///
/// - Complexity: O(log *n*), where *n* is the length of the queue.
@inlinable
public mutating func movedFirstAfterUpdating(with body: (inout Element) throws -> Void) rethrows -> Bool {
return try base.updateHeapLead(using: body, prioritizing: priority, by: areInIncreasingOrder) != nil
}
}
extension SomePriorityQueue where Base: RangeReplaceableCollection {
/// Creates an empty priority queue, with its ordering checked by the given
/// predicate, vending in the order of the given priority.
///
/// - Precondition: `areInIncreasingOrder` must be a strict weak ordering
/// over the elements.
///
/// - Parameter priority: Indicates which end of the ordering scale that
/// elements will be read from.
/// - Parameter areInIncreasingOrder: A predicate that returns `true` if its
/// first argument should be ordered before its second argument;
/// otherwise, `false`.
@inlinable
public init(priority: HeapPriority, by areInIncreasingOrder: @escaping (Element, Element) -> Bool) {
self.init(base: Base(), priority: priority, by: areInIncreasingOrder)
}
/// Creates a priority queue from the elements of the given sequence, with
/// its ordering checked by the given predicate, vending in the order of the
/// given priority.
///
/// - Precondition: `areInIncreasingOrder` must be a strict weak ordering
/// over the elements. `elements` must be finite.
///
/// - Parameter elements: The elements to form the initial queue.
/// - Parameter priority: Indicates which end of the ordering scale that
/// elements will be read from.
/// - Parameter areInIncreasingOrder: A predicate that returns `true` if its
/// first argument should be ordered before its second argument;
/// otherwise, `false`.
///
/// - Postcondition: The values from `elements` are rearranged to ease the
/// vending of either the lowest ranked (when `priority == .min`) or
/// highest ranked (when `priority == .max`) elements as determined by
/// `areInIncreasingOrder`.
///
/// - Complexity: At least O(*n*), where *n* is the length of `elements`.
@inlinable
public init<S: Sequence>(contentsOf elements: S, priority: HeapPriority, by areInIncreasingOrder: @escaping (Element, Element) -> Bool) where S.Element == Element {
self.init(base: Base(elements), priority: priority, by: areInIncreasingOrder)
}
/// Adds a given element to the queue.
///
/// If the queue does not have enough capacity to store another element,
/// additional storage is allocated before inserting `newElement`.
///
/// - Parameter newElement: The element to insert into the queue.
///
/// - Postcondition: `newElement` is part of the queue, but not necessarily
/// at `first`.
///
/// - Complexity: Amortized O(log *n*), where *n* is the length of the
/// queue.
@inlinable
public mutating func push(_ newElement: Element) {
base.append(newElement)
base.partitionIntoHeap(at: base.indices.last!, prioritizing: priority, by: areInIncreasingOrder)
}
/// Adds the elements of the given sequence to the queue.
///
/// If the queue does not have enough capacity to store more elements,
/// additional storage is allocated before inserting anything from
/// `newElements`.
///
/// - Parameter newElements: The elements to insert into the queue.
///
/// - Postcondition: The elements from `newElements` are part of the queue,
/// but none necessarily at `first`.
///
/// - Complexity: Amortized O(*m* × log *n*), where *m* is the length of
/// `newElements` and *n* is the length of the queue.
public mutating func push<S: Sequence>(contentsOf newElements: S) where S.Element == Element {
let originalCount = count
base.append(contentsOf: newElements)
var current = base.index(base.startIndex, offsetBy: originalCount)
let end = base.endIndex
while current < end {
current = base[...current].partitionIntoHeap(at: current, prioritizing: priority, by: areInIncreasingOrder)
}
}
/// Removes and returns the best-ranked element, if any.
///
/// - Returns: The lowest-ranked value when `priority` is `.min`, the
/// highest-ranked value when `priority` is `.max`, but `nil` when the
/// queue is empty.
///
/// - Postcondition: The returned best-ranked element is no longer part of
/// the queue. A second-best element is moved up to `first`.
///
/// - Complexity: O(log *n*), where *n* is the length of the queue.
@inlinable
public mutating func popFirst() -> Element? {
base.partitionOutHeapLead(prioritizing: priority, by: areInIncreasingOrder)
return base.popLast()
}
/// Removes and returns the best-ranked element.
///
/// - Precondition: `!isEmpty()`.
///
/// - Returns: The lowest-ranked value when `priority` is `.min`, the
/// highest-ranked value when `priority` is `.max`.
///
/// - Postcondition: The returned best-ranked element is no longer part of
/// the queue. A second-best element is moved up to `first`.
///
/// - Complexity: O(log *n*), where *n* is the length of the queue.
@inlinable
@discardableResult
public mutating func removeFirst() -> Element {
guard let oldFirst = popFirst() else {
preconditionFailure("Request first from empty priority queue")
}
return oldFirst
}
/// Removes the given count of best-ranked elements from the queue.
///
/// - Precondition: `0 ≤ k < count`.
///
/// - Parameter k: The number of elements to remove.
///
/// - Postcondition: The `k` best-ranked elements are no longer part of the
/// queue. The next-best element after all of those is moved up to
/// `first`.
///
/// - Complexity: O(*k* × log *n*), where *n* is the length of the queue.
public mutating func removeFirst(_ k: Int) {
precondition(0..<count ~= k)
var current = base.endIndex
let keptEnd = base.index(current, offsetBy: -k)
while current > keptEnd {
current = base[..<current].partitionOutHeapLead(prioritizing: priority, by: areInIncreasingOrder)
}
base.removeSubrange(current...)
}
/// Removes all elements from the queue.
///
/// - Parameter keepCapacity: Whether to retain the memory allocated for the
/// removed elements, in case new elements will be pushed in. If not
/// given, defaults to `false`.
///
/// - Postcondition: `isEmpty()`.
///
/// - Complexity: O(*n*), where *n* is the length of the queue.
@inlinable
public mutating func removeAll(keepingCapacity keepCapacity: Bool = false) {
base.removeAll(keepingCapacity: keepCapacity)
}
/// Prepares the queue to store the given number of elements, if possible.
///
/// - Parameter n: The requested number of elements to store.
///
/// - Postcondition: The underlying `Base` implementation of reserving
/// capacity is used.
@inlinable
public mutating func reserveCapacity(_ n: Int) {
base.reserveCapacity(n)
}
}
extension SomePriorityQueue: CustomDebugStringConvertible {
private var elementDescriptions: String {
let start = base.startIndex
guard let secondIndex = base.index(start, offsetBy: +1, limitedBy: base.endIndex) else { return "" }
var range = start..<secondIndex
var ranges = Array<Range<Base.Index>>()
repeat {
ranges.append(range)
let indexRange = base.indices[range]
let lowChildRange = base.indices.binaryHeapChildrenRange(of: indexRange.first!)
let highChildRange = base.indices.binaryHeapChildrenRange(of: indexRange.last!)
range = lowChildRange.lowerBound ..< highChildRange.upperBound
} while !range.isEmpty
return ranges.lazy.map { self.base[$0].lazy.map(String.init(describing:)).joined(separator: ", ") }.joined(separator: " | ")
}
public var debugDescription: String {
String(describing: Self.self) + "(" + elementDescriptions + ")"
}
}
/// A container controlling access to only its best-ranked element.
public typealias PriorityQueue<T> = SomePriorityQueue<Array<T>>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment