Skip to content

Instantly share code, notes, and snippets.

@khanlou
Last active July 28, 2020 16:54
Show Gist options
  • Save khanlou/770c24d5141e52e117642c4b03498966 to your computer and use it in GitHub Desktop.
Save khanlou/770c24d5141e52e117642c4b03498966 to your computer and use it in GitHub Desktop.
Smallest N items in a sequence, with binary search improvements
import Foundation
extension Array {
func insertionIndex(of element: Element, by areInIncreasingOrder: (Element, Element) -> Bool) -> Index? {
if isEmpty { return nil }
if let last = self.last, areInIncreasingOrder(last, element) { return nil }
var start = startIndex
var end = endIndex
while start < end {
let middle = start / 2 + end / 2
if areInIncreasingOrder(self[middle], element) {
start = middle + 1
} else if areInIncreasingOrder(element, self[middle]) {
end = middle
} else {
return middle
}
}
return start
}
}
extension Sequence {
func smallest(_ m: Int, usingTest areInIncreasingOrder: (Element, Element) -> Bool) -> [Element] {
var result = self.prefix(n).sorted(by: areInIncreasingOrder)
for e in self.dropFirst(n) {
if let insertionIndex = result.insertionIndex(of: e, by: areInIncreasingOrder) {
result.insert(e, at: insertionIndex)
result.removeLast()
}
}
return result
}
}
let start = Date()
let shuff = (0..<55000).shuffled()
let a = shuff.smallest(20, usingTest: <)
let end = Date()
print(a)
print(end.timeIntervalSince(start))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment