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