| 
 | 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 | +//===----------------------------------------------------------------------===//  | 
 | 11 | + | 
 | 12 | +extension Collection {  | 
 | 13 | +  /// Returns the first k elements of this collection when it's sorted using  | 
 | 14 | +  /// the given predicate as the comparison between elements.  | 
 | 15 | +  ///  | 
 | 16 | +  /// This example partially sorts an array of integers to retrieve its three  | 
 | 17 | +  /// smallest values:  | 
 | 18 | +  ///  | 
 | 19 | +  ///     let numbers = [7,1,6,2,8,3,9]  | 
 | 20 | +  ///     let smallestThree = numbers.sortedPrefix(3, by: <)  | 
 | 21 | +  ///     // [1, 2, 3]  | 
 | 22 | +  ///  | 
 | 23 | +  /// If you need to sort a collection but only need access to a prefix of its  | 
 | 24 | +  /// elements, using this method can give you a performance boost over sorting  | 
 | 25 | +  /// the entire collection. The order of equal elements is guaranteed to be  | 
 | 26 | +  /// preserved.  | 
 | 27 | +  ///  | 
 | 28 | +  /// - Parameter count: The k number of elements to prefix.  | 
 | 29 | +  /// - Parameter areInIncreasingOrder: A predicate that returns true if its  | 
 | 30 | +  /// first argument should be ordered before its second argument;  | 
 | 31 | +  /// otherwise, false.  | 
 | 32 | +  ///  | 
 | 33 | +  /// - Complexity: O(k log k + nk)  | 
 | 34 | +  public func sortedPrefix(  | 
 | 35 | +    _ count: Int,  | 
 | 36 | +    by areInIncreasingOrder: (Element, Element) throws -> Bool  | 
 | 37 | +  ) rethrows -> [Self.Element] {  | 
 | 38 | +    assert(count >= 0, """  | 
 | 39 | +      Cannot prefix with a negative amount of elements!  | 
 | 40 | +      """  | 
 | 41 | +    )  | 
 | 42 | + | 
 | 43 | +    // Do nothing if we're prefixing nothing.  | 
 | 44 | +    guard count > 0 else {  | 
 | 45 | +      return []  | 
 | 46 | +    }  | 
 | 47 | + | 
 | 48 | +    // Make sure we are within bounds.  | 
 | 49 | +    let prefixCount = Swift.min(count, self.count)  | 
 | 50 | + | 
 | 51 | +    // If we're attempting to prefix more than 10% of the collection, it's  | 
 | 52 | +    // faster to sort everything.  | 
 | 53 | +    guard prefixCount < (self.count / 10) else {  | 
 | 54 | +      return Array(try sorted(by: areInIncreasingOrder).prefix(prefixCount))  | 
 | 55 | +    }  | 
 | 56 | + | 
 | 57 | +    var result = try self.prefix(prefixCount).sorted(by: areInIncreasingOrder)  | 
 | 58 | +    for e in self.dropFirst(prefixCount) {  | 
 | 59 | +      if let last = result.last, try areInIncreasingOrder(last, e) {  | 
 | 60 | +        continue  | 
 | 61 | +      }  | 
 | 62 | +      let insertionIndex =  | 
 | 63 | +        try result.partitioningIndex { try areInIncreasingOrder(e, $0) }  | 
 | 64 | +      let isLastElement = insertionIndex == result.endIndex  | 
 | 65 | +      result.removeLast()  | 
 | 66 | +      if isLastElement {  | 
 | 67 | +        result.append(e)  | 
 | 68 | +      } else {  | 
 | 69 | +        result.insert(e, at: insertionIndex)  | 
 | 70 | +      }  | 
 | 71 | +    }  | 
 | 72 | + | 
 | 73 | +    return result  | 
 | 74 | +  }  | 
 | 75 | +}  | 
 | 76 | + | 
 | 77 | +extension Collection where Element: Comparable {  | 
 | 78 | +  /// Returns the first k elements of this collection when it's sorted in  | 
 | 79 | +  /// ascending order.  | 
 | 80 | +  ///  | 
 | 81 | +  /// This example partially sorts an array of integers to retrieve its three  | 
 | 82 | +  /// smallest values:  | 
 | 83 | +  ///  | 
 | 84 | +  ///     let numbers = [7,1,6,2,8,3,9]  | 
 | 85 | +  ///     let smallestThree = numbers.sortedPrefix(3)  | 
 | 86 | +  ///     // [1, 2, 3]  | 
 | 87 | +  ///  | 
 | 88 | +  /// If you need to sort a collection but only need access to a prefix of its  | 
 | 89 | +  /// elements, using this method can give you a performance boost over sorting  | 
 | 90 | +  /// the entire collection. The order of equal elements is guaranteed to be  | 
 | 91 | +  /// preserved.  | 
 | 92 | +  ///  | 
 | 93 | +  /// - Parameter count: The k number of elements to prefix.  | 
 | 94 | +  ///  | 
 | 95 | +  /// - Complexity: O(k log k + nk)  | 
 | 96 | +  public func sortedPrefix(_ count: Int) -> [Element] {  | 
 | 97 | +    return sortedPrefix(count, by: <)  | 
 | 98 | +  }  | 
 | 99 | +}  | 
0 commit comments