Skip to content

Commit 4362197

Browse files
rockbrunorockbruno
committed
Add in place partial sorting
1 parent 5429d3b commit 4362197

File tree

1 file changed

+182
-64
lines changed

1 file changed

+182
-64
lines changed

‎Sources/Algorithms/PartialSort.swift

Lines changed: 182 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -1,89 +1,207 @@
1-
import Foundation
1+
//===----------------------------------------------------------------------===//
2+
//
3+
// This source file is part of the Swift Algorithms open source project
4+
//
5+
// Copyright (c) 2020 Apple Inc. and the Swift project authors
6+
// Licensed under Apache License v2.0 with Runtime Library Exception
7+
//
8+
// See https://swift.org/LICENSE.txt for license information
9+
//
10+
//===----------------------------------------------------------------------===//
211

3-
private final class Heap<T> {
4-
5-
typealias Comparator = (T,T) -> Bool
6-
7-
private var elements: [T]
8-
private let priority: Comparator
9-
10-
init<S: Sequence>(elements: S, priority: @escaping Comparator) where S.Element == T {
11-
self.priority = priority
12-
self.elements = Array(elements)
13-
if elements.isEmpty == false {
14-
for i in stride(from: (count / 2) - 1, to: -1, by: -1) {
15-
siftDown(i)
16-
}
17-
}
12+
extension Sequence {
13+
/// Returns the elements of the sequence such that the 0...k range contains
14+
/// the first k sorted elements in this sequence, using the given predicate
15+
/// as the comparison between elements.
16+
///
17+
/// This example partially sorts an array of integers to retrieve its three
18+
/// smallest values:
19+
///
20+
/// let numbers = [7,1,6,2,8,3,9]
21+
/// let almostSorted = numbers.partiallySorted(3, <)
22+
/// // [1, 2, 3, 9, 7, 6, 8]
23+
/// let smallestThree = almostSorted.prefix(3)
24+
/// // [1, 2, 3]
25+
///
26+
/// The order of equal elements is not guaranteed to be preserved, and the
27+
/// order of the remaining elements is unspecified.
28+
///
29+
/// If you need to sort a sequence but only need access to a prefix of its elements,
30+
/// using this method can give you a performance boost over sorting the entire
31+
/// sequence.
32+
///
33+
/// - Parameter count: The k number of elements to partially sort.
34+
/// - Parameter areInIncreasingOrder: A predicate that returns true if its first argument should
35+
/// be ordered before its second argument; otherwise, false.
36+
///
37+
/// - Complexity: O(k log n)
38+
public func partiallySorted(_ count: Int, by areInIncreasingOrder: (Element, Element) throws -> Bool) rethrows -> [Element] {
39+
var result = ContiguousArray(self)
40+
try result.partiallySort(count, by: areInIncreasingOrder)
41+
return Array(result)
1842
}
43+
}
1944

20-
private func leftChildIndex(of index: Int) -> Int {
21-
return (2 * index) + 1
45+
extension Sequence where Element: Comparable {
46+
/// Returns the elements of the sequence such that the 0...k range contains
47+
/// the first k smallest elements in this sequence.
48+
///
49+
/// This example partially sorts an array of integers to retrieve its three
50+
/// smallest values:
51+
///
52+
/// let numbers = [7,1,6,2,8,3,9]
53+
/// let almostSorted = numbers.partiallySorted(3)
54+
/// // [1, 2, 3, 9, 7, 6, 8]
55+
/// let smallestThree = almostSorted.prefix(3)
56+
/// // [1, 2, 3]
57+
///
58+
/// The order of equal elements is not guaranteed to be preserved, and the
59+
/// order of the remaining elements is unspecified.
60+
///
61+
/// If you need to sort a sequence but only need access to a prefix of its elements,
62+
/// using this method can give you a performance boost over sorting the entire
63+
/// sequence.
64+
///
65+
/// - Parameter count: The k number of elements to partially sort, in ascending order.
66+
///
67+
/// - Complexity: O(k log n)
68+
public func partiallySorted(_ count: Int) -> [Element] {
69+
return partiallySorted(count, by: <)
2270
}
71+
}
2372

24-
private func rightChild(of index: Int) -> Int {
25-
return (2 * index) + 2
73+
extension MutableCollection where Self: RandomAccessCollection, Index == Int {
74+
/// Rearranges this collection such that the 0...k range contains the first
75+
/// k sorted elements in this collection, using the given predicate as the
76+
/// comparison between elements.
77+
///
78+
/// This example partially sorts an array of integers to retrieve its three
79+
/// smallest values:
80+
///
81+
/// var numbers = [7,1,6,2,8,3,9]
82+
/// numbers.partiallySort(3, <)
83+
/// // [1, 2, 3, 9, 7, 6, 8]
84+
/// let smallestThree = numbers.prefix(3)
85+
/// // [1, 2, 3]
86+
///
87+
/// The order of equal elements is not guaranteed to be preserved, and the
88+
/// order of the remaining elements is unspecified.
89+
///
90+
/// If you need to sort a collection but only need access to a prefix of its
91+
/// elements, using this method can give you a performance boost over sorting
92+
/// the entire collection.
93+
///
94+
/// - Parameter count: The k number of elements to partially sort.
95+
/// - Parameter areInIncreasingOrder: A predicate that returns true if its first argument should
96+
/// be ordered before its second argument; otherwise, false.
97+
///
98+
/// - Complexity: O(k log n)
99+
public mutating func partiallySort(_ count: Int, by areInIncreasingOrder: (Element, Element) throws -> Bool) rethrows {
100+
try __partiallySort(count, by: areInIncreasingOrder)
26101
}
102+
}
27103

28-
private func parentIndex(of index: Int) -> Int {
29-
return (index - 1) / 2
104+
extension MutableCollection where Self: RandomAccessCollection, Element: Comparable, Index == Int {
105+
/// Rearranges this collection such that the 0...k range contains the first
106+
/// k smallest elements in this collection.
107+
///
108+
/// This example partially sorts an array of integers to retrieve its three
109+
/// smallest values:
110+
///
111+
/// var numbers = [7,1,6,2,8,3,9]
112+
/// numbers.partiallySort(3)
113+
/// // [1, 2, 3, 9, 7, 6, 8]
114+
/// let smallestThree = numbers.prefix(3)
115+
/// // [1, 2, 3]
116+
///
117+
/// The order of equal elements is not guaranteed to be preserved, and the
118+
/// order of the remaining elements is unspecified.
119+
///
120+
/// If you need to sort a collection but only need access to a prefix of its
121+
/// elements, using this method can give you a performance boost over sorting
122+
/// the entire collection.
123+
///
124+
/// - Parameter count: The k number of elements to partially sort, in ascending order.
125+
///
126+
/// - Complexity: O(k log n)
127+
public mutating func partiallySort(_ count: Int) {
128+
partiallySort(count, by: <)
30129
}
130+
}
31131

32-
private func isHigherPriority(_ a: Int, _ b: Int) -> Bool {
33-
return priority(elements[a], elements[b])
34-
}
132+
//===----------------------------------------------------------------------===//
133+
// __partiallySort(_:by:)
134+
//===----------------------------------------------------------------------===//
35135

36-
private func highestPriorityIndex(of index: Int) -> Int {
37-
let left = highestPriorityIndex(of: index, and: leftChildIndex(of: index))
38-
let right = highestPriorityIndex(of: index, and: rightChild(of: index))
39-
return highestPriorityIndex(of: left, and: right)
40-
}
136+
extension MutableCollection where Self: RandomAccessCollection, Index == Int {
137+
typealias Priority = (Element, Element) throws -> Bool
41138

42-
private func highestPriorityIndex(of parent: Int, and child: Int) -> Int {
43-
guard child < elements.count else {
44-
return parent
139+
/// Partially sorts this array by using an in place heapsort that stops after we find the desired k amount
140+
/// of elements. The heap is stored and processed in reverse order so that the array doesn't have to be flipped
141+
/// once the final result is found.
142+
///
143+
/// Complexity: O(k log n)
144+
mutating func __partiallySort(_ k: Int, by areInIncreasingOrder: Priority) rethrows {
145+
assert(k >= 0, "Attempted to partially sort with a negative amount of elements!")
146+
assert(k <= count,
147+
"Attempted to partially sort an amount of elements larger than this Sequence's size!")
148+
guard isEmpty == false else {
149+
return
45150
}
46-
guard isHigherPriority(child, parent) else {
47-
return parent
151+
var heapEndIndex = 0
152+
for i in ((count / 2) + 1)..<count {
153+
try siftDown(i, by: areInIncreasingOrder, heapEndIndex: heapEndIndex)
48154
}
49-
return child
50-
}
51-
52-
func dequeue() -> T? {
53-
guard elements.count > 0 else {
54-
return nil
155+
var iterator = (0..<k).makeIterator()
156+
_ = iterator.next()
157+
swapAt(count - 1, heapEndIndex)
158+
heapEndIndex += 1
159+
while let _ = iterator.next() {
160+
try siftDown(count - 1, by: areInIncreasingOrder, heapEndIndex: heapEndIndex)
161+
swapAt(count - 1, heapEndIndex)
162+
heapEndIndex += 1
55163
}
56-
elements.swapAt(0, elements.count - 1)
57-
let element = elements.popLast()
58-
siftDown(0)
59-
return element
60164
}
61165

62-
private func siftDown(_ i: Int) {
63-
let indexToSwap = highestPriorityIndex(of: i)
166+
/// Sifts down an element from this heap.
167+
/// The heap is stored in reverse order, so sifting down will actually move the element up in the heap array.
168+
///
169+
/// - Parameter i: The element index to sift down
170+
/// - Parameter by: The predicate to use when determining the priority of elements in the heap
171+
/// - Parameter heapEndIndex: The index, in reverse order, where the heap ends.
172+
private mutating func siftDown(_ i: Int, by priority: Priority, heapEndIndex: Int) rethrows {
173+
let indexToSwap = try highestPriorityIndex(of: i, by: priority, heapEndIndex: heapEndIndex)
64174
guard indexToSwap != i else {
65175
return
66176
}
67-
elements.swapAt(indexToSwap, i)
68-
siftDown(indexToSwap)
177+
swapAt(i, indexToSwap)
178+
try siftDown(indexToSwap, by: priority, heapEndIndex: heapEndIndex)
69179
}
70-
}
71180

72-
extension Collection {
73-
func partiallySorted(_ count: Int, by: @escaping (Element, Element) -> Bool) -> [Element] {
74-
assert(count >= 0 && count < self.count, "Are you crazy?")
75-
let heap = Heap<Element>(elements: self, priority: by)
76-
return [Element](unsafeUninitializedCapacity: count) { buffer, initializedCount in
77-
for i in 0..<count {
78-
buffer[i] = heap.dequeue()!
79-
}
80-
initializedCount = count
81-
}
181+
private func highestPriorityIndex(of index: Int, by priority: Priority, heapEndIndex: Int) rethrows -> Int {
182+
let reverseHeapTrueIndex = self.count - 1 - index
183+
let leftChild = index - (leftChildIndex(of: reverseHeapTrueIndex) - reverseHeapTrueIndex)
184+
let rightChild = index - (rightChildIndex(of: reverseHeapTrueIndex) - reverseHeapTrueIndex)
185+
let left = try highestPriorityIndex(of: index, and: leftChild, by: priority, heapEndIndex: heapEndIndex)
186+
let right = try highestPriorityIndex(of: index, and: rightChild, by: priority, heapEndIndex: heapEndIndex)
187+
return try highestPriorityIndex(of: left, and: right, by: priority, heapEndIndex: heapEndIndex)
82188
}
83-
}
84189

85-
extension Collection where Element: Comparable {
86-
func partiallySorted(_ count: Int) -> [Element] {
87-
return partiallySorted(count, by: <)
190+
private func leftChildIndex(of index: Int) -> Int {
191+
return (2 * index) + 1
192+
}
193+
194+
private func rightChildIndex(of index: Int) -> Int {
195+
return (2 * index) + 2
196+
}
197+
198+
private func highestPriorityIndex(of parent: Int, and child: Int, by priority: Priority, heapEndIndex: Int) rethrows -> Int {
199+
guard child >= heapEndIndex else {
200+
return parent
201+
}
202+
guard try priority(self[child], self[parent]) else {
203+
return parent
204+
}
205+
return child
88206
}
89207
}

0 commit comments

Comments
 (0)