Skip to content

Instantly share code, notes, and snippets.

@timvermeulen
Created January 8, 2019 21:37
Show Gist options
  • Save timvermeulen/5cf117221157fd23d6e083e2c1f19aa2 to your computer and use it in GitHub Desktop.
Save timvermeulen/5cf117221157fd23d6e083e2c1f19aa2 to your computer and use it in GitHub Desktop.
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()
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment