Partition.swift 1.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051
  1. //===----------------------------------------------------------------------===//
  2. //
  3. // This source file is part of the Swift Algorithms open source project
  4. //
  5. // Copyright (c) 2020-2021 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. // partitioningIndex(where:)
  13. //===----------------------------------------------------------------------===//
  14. extension Collection {
  15. /// Returns the start index of the partition of a collection that matches
  16. /// the given predicate.
  17. ///
  18. /// The collection must already be partitioned according to the predicate.
  19. /// That is, there should be an index `i` where for every element in
  20. /// `collection[..<i]` the predicate is `false`, and for every element in
  21. /// `collection[i...]` the predicate is `true`.
  22. ///
  23. /// - Parameter belongsInSecondPartition: A predicate that partitions the
  24. /// collection.
  25. /// - Returns: The index of the first element in the collection for which
  26. /// `predicate` returns `true`, or `endIndex` if there are no elements
  27. /// for which `predicate` returns `true`.
  28. ///
  29. /// - Complexity: O(log *n*), where *n* is the length of this collection if
  30. /// the collection conforms to `RandomAccessCollection`, otherwise O(*n*).
  31. func partitioningIndex(
  32. where belongsInSecondPartition: (Element) throws -> Bool
  33. ) rethrows -> Index {
  34. var n = count
  35. var l = startIndex
  36. while n > 0 {
  37. let half = n / 2
  38. let mid = index(l, offsetBy: half)
  39. if try belongsInSecondPartition(self[mid]) {
  40. n = half
  41. } else {
  42. l = index(after: mid)
  43. n -= half + 1
  44. }
  45. }
  46. return l
  47. }
  48. }