Guide docs

Guides/PartialSort.md

@@ -0,0 +1,59 @@
+# Partial Sort
+
+[[Source](https://github.com/apple/swift-algorithms/blob/main/Sources/Algorithms/PartialSort.swift) | 
+ [Tests](https://github.com/apple/swift-algorithms/blob/main/Tests/SwiftAlgorithmsTests/PartialSortTests.swift)]
+
+Returns a collection such that the `0...k` range contains the first k sorted elements of a sequence. 
+The order of equal elements is not guaranteed to be preserved, and the order of the remaining elements is unspecified.
+
+If you need to sort a sequence but only need access to a prefix of its elements, 
+using this method can give you a performance boost over sorting the entire sequence.
+
+```swift
+let numbers = [7,1,6,2,8,3,9]
+let almostSorted = numbers.partiallySorted(3, <)
+// [1, 2, 3, 9, 7, 6, 8]
+```
+
+## Detailed Design
+
+This adds the in place `MutableCollection` method shown below:
+
+```swift
+extension Sequence {
+    func partiallySort(_ count: Int, by areInIncreasingOrder: (Element, Element) throws -> Bool) rethrows
+}
+```
+
+Additionally, versions of this method that return a new array and abstractions for `Comparable` types are also provided:
+
+```swift
+extension MutableCollection where Self: RandomAccessCollection, Element: Comparable {
+    public mutating func partiallySort(_ count: Int)
+}
+
+extension Sequence {
+    public func partiallySorted(_ count: Int, by areInIncreasingOrder: (Element, Element) throws -> Bool) rethrows -> [Element]
+}
+
+extension Sequence where Element: Comparable {
+    public func partiallySorted(_ count: Int) -> [Element]
+}
+```
+
+### Complexity
+
+Partially sorting is a O(_k log n_) operation, where _k_ is the number of elements to sort
+and _n_ is the length of the sequence.
+
+`partiallySort(_:by:)` is a slight generalization of a priority queue. It's implemented
+as an in line heapsort that stops after _k_ runs.
+
+### Comparison with other languages
+
+**C++:** The `<algorithm>` library defines a `partial_sort` function with similar
+semantics to this one.
+
+**Python:** Defines a `heapq` priority queue that can be used to manually 
+achieve the same result.
+

README.md

@@ -28,6 +28,10 @@ Read more about the package, and the intent behind it, in the [announcement on s
 - [`randomStableSample(count:)`, `randomStableSample(count:using:)`](https://github.com/apple/swift-algorithms/blob/main/Guides/RandomSampling.md): Randomly selects a specific number of elements from a collection, preserving their original relative order.
 - [`uniqued()`, `uniqued(on:)`](https://github.com/apple/swift-algorithms/blob/main/Guides/Unique.md): The unique elements of a collection, preserving their order.
 
+#### Partial sorting
+
+- [`partiallySorted(_:by:)`](https://github.com/apple/swift-algorithms/blob/main/Guides/PartialSort.md): Sorts a sequence only up to a specific index, leaving the remaining elements unsorted.
+
 #### Other useful operations
 
 - [`chunked(by:)`, `chunked(on:)`](https://github.com/apple/swift-algorithms/blob/main/Guides/Chunked.md): Eager and lazy operations that break a collection into chunks based on either a binary predicate or when the result of a projection changes.

Sources/Algorithms/PartialSort.swift

@@ -70,7 +70,7 @@ extension Sequence where Element: Comparable {
     }
 }
 
-extension MutableCollection where Self: RandomAccessCollection, Index == Int {
+extension MutableCollection where Self: RandomAccessCollection {
     /// Rearranges this collection such that the 0...k range contains the first
     /// k sorted elements in this collection, using the given predicate as the
     /// comparison between elements.
@@ -101,7 +101,7 @@ extension MutableCollection where Self: RandomAccessCollection, Index == Int {
     }
 }
 
-extension MutableCollection where Self: RandomAccessCollection, Element: Comparable, Index == Int {
+extension MutableCollection where Self: RandomAccessCollection, Element: Comparable {
     /// Rearranges this collection such that the 0...k range contains the first
     /// k smallest elements in this collection.
     ///
@@ -133,11 +133,11 @@ extension MutableCollection where Self: RandomAccessCollection, Element: Compara
 // __partiallySort(_:by:)
 //===----------------------------------------------------------------------===//
 
-extension MutableCollection where Self: RandomAccessCollection, Index == Int {
+extension MutableCollection where Self: RandomAccessCollection {
     typealias Priority = (Element, Element) throws -> Bool
 
-    /// Partially sorts this array by using an in place heapsort that stops after we find the desired k amount
-    /// of elements. The heap is stored and processed in reverse order so that the array doesn't have to be flipped
+    /// Partially sorts this collection by using an in place heapsort that stops after we find the desired k amount
+    /// of elements. The heap is stored and processed in reverse order so that the collection doesn't have to be flipped
     /// once the final result is found.
     ///
     /// Complexity: O(k log n)
@@ -154,17 +154,17 @@ extension MutableCollection where Self: RandomAccessCollection, Index == Int {
         }
         var iterator = (0..<k).makeIterator()
         _ = iterator.next()
-        swapAt(count - 1, heapEndIndex)
+        swapAt(index(before: endIndex), index(startIndex, offsetBy: heapEndIndex))
         heapEndIndex += 1
         while let _ = iterator.next() {
             try siftDown(count - 1, by: areInIncreasingOrder, heapEndIndex: heapEndIndex)
-            swapAt(count - 1, heapEndIndex)
+            swapAt(index(before: endIndex), index(startIndex, offsetBy: heapEndIndex))
             heapEndIndex += 1
         }
     }
 
     /// Sifts down an element from this heap.
-    /// The heap is stored in reverse order, so sifting down will actually move the element up in the heap array.
+    /// The heap is stored in reverse order, so sifting down will actually move the element up in the heap.
     ///
     /// - Parameter i: The element index to sift down
     /// - Parameter by: The predicate to use when determining the priority of elements in the heap
@@ -174,7 +174,7 @@ extension MutableCollection where Self: RandomAccessCollection, Index == Int {
         guard indexToSwap != i else {
             return
         }
-        swapAt(i, indexToSwap)
+        swapAt(index(startIndex, offsetBy: i), index(startIndex, offsetBy: indexToSwap))
         try siftDown(indexToSwap, by: priority, heapEndIndex: heapEndIndex)
     }
 
@@ -199,7 +199,9 @@ extension MutableCollection where Self: RandomAccessCollection, Index == Int {
         guard child >= heapEndIndex else {
             return parent
         }
-        guard try priority(self[child], self[parent]) else {
+        let childElement = self[index(startIndex, offsetBy: child)]
+        let parentElement = self[index(startIndex, offsetBy: parent)]
+        guard try priority(childElement, parentElement) else {
             return parent
         }
         return child