heapsort4.swift

/// Complexity: O(log n)
public mutating func remove() -> Element? {
    guard !isEmpty else { return nil }
    elements.swapAt(0, count - 1)
    defer {
        shiftDown(from: 0)
    }
    return elements.removeLast()
}

private mutating func shiftDown(from index: Int) {
    var parent = index
    while true {
        let leftChild = leftChildIndex(forParentAt: parent)
        let rightChild = rightChildIndex(forParentAt: parent)

        var next = parent

        if leftChild < count && sort(elements[leftChild], elements[next]) {
            next = leftChild
        }
        if rightChild < count && sort(elements[rightChild], elements[next]) {
            next = rightChild
        }

        if next == parent {
            // No need to shift down any more
            break
        }

        elements.swapAt(parent, next)
        parent = next
    }
}