This can be further optimized with a custom heap implementation. Once there are k elements in the heap, we know that whenever we remove an element from the hea

This can be further optimized with a custom heap implementation. Once there are k elements in the heap, we know that whenever we remove an element from the heap, we will immediately be adding a new one. In a traditional heap implementation, popping is performed by replacing the root with the bottom-right-most element and heapifying down. This is potentially 2 * log(k) comparisons. Pushing is performed by adding the new element to the bottom-right and heapifying up, which is up to log(k) comparisons. We can combine the pop and push into a single operation and reduce the 3 * log(k) comparisons to just the 2 * log(k) comparisons. Instead of replacing the root with the bottom-right-most element, just replace it with the new element and heapify down. This not only reduces the worst-case number of comparison; it improves the expected number of comparisons since the bottom-right element is likely to still belong at the bottom and will be pushed down the entire tree, while the new element is more likely to belong somewhere shallower.