Optimize

Author: rockbruno

Sources/Algorithms/PartialSort.swift

@@ -50,11 +50,12 @@ extension Collection {
 
     var result = try self.prefix(count).sorted(by: areInIncreasingOrder)
     for e in self.dropFirst(count) {
-      if let last = result.last, try areInIncreasingOrder(last, e) { continue }
-      if let insertionIndex = try result.firstIndex(where: { try areInIncreasingOrder(e, $0) }) {
-        result.insert(e, at: insertionIndex)
-        result.removeLast()
+      if let last = result.last, try areInIncreasingOrder(last, e) {
+        continue
       }
+      let insertionIndex = try result.partitioningIndex { try areInIncreasingOrder(e, $0) }
+      result.removeLast()
+      result.insert(e, at: insertionIndex)
     }
 
     return result

Tests/SwiftAlgorithmsTests/PartialSortTests.swift

@@ -67,6 +67,31 @@ final class PartialSortTests: XCTestCase {
       [1, 2, 3, 4, 7, 20, 70, 90, 100].sortedPrefix(5, by: <),
       [1, 2, 3, 4, 7]
     )
+
+    XCTAssertEqual(
+      [4, 5, 6, 1, 2, 3].sortedPrefix(3, by: <),
+      [1, 2, 3]
+    )
+
+    XCTAssertEqual(
+      [4, 5, 9, 8, 7, 6].sortedPrefix(3, by: <),
+      [4, 5, 6]
+    )
+
+    XCTAssertEqual(
+      [4, 3, 2, 1].sortedPrefix(1, by: <),
+      [1]
+    )
+
+    XCTAssertEqual(
+      [4, 2, 1, 3].sortedPrefix(3, by: >),
+      [4, 3, 2]
+    )
+
+    XCTAssertEqual(
+      [4, 2, 1, 3].sortedPrefix(3, by: <),
+      [1, 2, 3]
+    )
   }
 
   func testSortedPrefixComparable() {