NonEmptyMaxHeap.swift

struct NonEmptyMaxHeap<Element: Comparable> {
    private(set) var elements: [Element]
    
    init(root: Element) {
        self.elements = [root]
    }
}

extension NonEmptyMaxHeap {
    mutating func insert(_ element: Element) {
        var index = elements.count
        elements.append(element)
        
        while index > 0 {
            let parentIndex = (index - 1) / 2
            let parent = elements[parentIndex]
            
            guard parent < element else { break }
            
            elements[index] = parent
            index = parentIndex
        }
        
        elements[index] = element
    }
    
    var root: Element {
        return elements.first!
    }
    
    mutating func replaceRoot(with element: Element) {
        var index = 0
        
        while case var childIndex = 2 * index + 1, childIndex < elements.count {
            var child = elements[childIndex]
            let rightIndex = childIndex + 1
            
            if rightIndex < elements.count {
                let right = elements[rightIndex]
                
                if child < right {
                    child = right
                    childIndex = rightIndex
                }
            }
            
            guard element < child else { break }
            
            elements[index] = child
            index = childIndex
        }
        
        elements[index] = element
    }
}

extension Sequence where Element: Comparable {
    func min(_ n: Int) -> [Element] {
        var iterator = makeIterator()
        guard let first = iterator.next() else { return [] }
        
        var heap = NonEmptyMaxHeap(root: first)
        var heapSize = 1
        
        while heapSize < n, let element = iterator.next() {
            heap.insert(element)
            heapSize += 1
        }
        
        while let element = iterator.next() {
            if element < heap.root {
                heap.replaceRoot(with: element)
            }
        }
        
        return heap.elements.sorted()
    }
}