All combinations (#51)

Add an overload of combinations(ofCount:) that accepts a range.

Author: mdznr

Guides/Combinations.md

@@ -36,6 +36,27 @@ for combo in numbers2.combinations(ofCount: 2) {
 // [10, 10]
 ```
 
+Given a range, the `combinations(ofCount:)` method returns a sequence of all
+the different combinations of the given sizes of a collection’s elements in
+increasing order of size.
+
+```swift
+let numbers = [10, 20, 30, 40]
+for combo in numbers.combinations(ofCount: 2...3) {
+    print(combo)
+}
+// [10, 20]
+// [10, 30]
+// [10, 40]
+// [20, 30]
+// [20, 40]
+// [30, 40]
+// [10, 20, 30]
+// [10, 20, 40]
+// [10, 30, 40]
+// [20, 30, 40]
+```
+
 ## Detailed Design
 
 The `combinations(ofCount:)` method is declared as a  `Collection` extension,
@@ -56,9 +77,9 @@ array at every index advancement. `Combinations` does conform to
 
 ### Complexity
 
-Calling `combinations(ofCount:)` accesses the count of the collection, so it’s an
-O(1) operation for random-access collections, or an O(_n_) operation otherwise.
-Creating the iterator for a `Combinations` instance and each call to
+Calling `combinations(ofCount:)` accesses the count of the collection, so it’s
+an O(1) operation for random-access collections, or an O(_n_) operation
+otherwise. Creating the iterator for a `Combinations` instance and each call to
 `Combinations.Iterator.next()` is an O(_n_) operation.
 
 ### Naming

README.md

@@ -8,7 +8,7 @@ Read more about the package, and the intent behind it, in the [announcement on s
 
 #### Combinations / permutations
 
-- [`combinations(ofCount:)`](https://github.com/apple/swift-algorithms/blob/main/Guides/Combinations.md): Combinations of a particular size of the elements in a collection.
+- [`combinations(ofCount:)`](https://github.com/apple/swift-algorithms/blob/main/Guides/Combinations.md): Combinations of particular sizes of the elements in a collection.
 - [`permutations(ofCount:)`](https://github.com/apple/swift-algorithms/blob/main/Guides/Permutations.md): Permutations of a particular size of the elements in a collection, or of the full collection.
 
 #### Mutating algorithms

Sources/Algorithms/Combinations.swift

@@ -2,7 +2,7 @@
 //
 // This source file is part of the Swift Algorithms open source project
 //
-// Copyright (c) 2020 Apple Inc. and the Swift project authors
+// Copyright (c) 2020-2021 Apple Inc. and the Swift project authors
 // Licensed under Apache License v2.0 with Runtime Library Exception
 //
 // See https://swift.org/LICENSE.txt for license information
@@ -15,16 +15,51 @@ public struct Combinations<Base: Collection> {
   public let base: Base
 
   @usableFromInline
-  internal var k: Int
+  internal let baseCount: Int
 
+  /// The range of accepted sizes of combinations.
+  /// - Note: This may be `nil` if the attempted range entirely exceeds the
+  /// upper bounds of the size of the `base` collection.
+  @usableFromInline
+  internal let kRange: Range<Int>?
+  
+  /// Initializes a `Combinations` for all combinations of `base` of size `k`.
+  /// - Parameters:
+  ///   - base: The collection to iterate over for combinations.
+  ///   - k: The expected size of each combination.
   @usableFromInline
   internal init(_ base: Base, k: Int) {
+    self.init(base, kRange: k...k)
+  }
+  
+  /// Initializes a `Combinations` for all combinations of `base` of sizes
+  /// within a given range.
+  /// - Parameters:
+  ///   - base: The collection to iterate over for combinations.
+  ///   - kRange: The range of accepted sizes of combinations.
+  @usableFromInline
+  internal init<R: RangeExpression>(
+    _ base: Base, kRange: R
+  ) where R.Bound == Int {
+    let range = kRange.relative(to: 0 ..< .max)
     self.base = base
-    self.k = base.count < k ? -1 : k
+    let baseCount = base.count
+    self.baseCount = baseCount
+    let upperBound = baseCount + 1
+    self.kRange = range.lowerBound < upperBound
+      ? range.clamped(to: 0 ..< upperBound)
+      : nil
   }
-
+  
+  /// The total number of combinations.
   @inlinable
   public var count: Int {
+    guard let k = self.kRange else { return 0 }
+    let n = baseCount
+    if k == 0 ..< (n + 1) {
+      return 1 << n
+    }
+    
     func binomial(n: Int, k: Int) -> Int {
       switch k {
       case n, 0: return 1
@@ -34,9 +69,9 @@ public struct Combinations<Base: Collection> {
       }
     }
 
-    return k >= 0
-      ? binomial(n: base.count, k: k)
-      : 0
+    return k.map {
+      binomial(n: n, k: $0)
+    }.reduce(0, +)
   }
 }
 
@@ -46,18 +81,26 @@ extension Combinations: Sequence {
     @usableFromInline
     internal let base: Base
 
+    /// The current range of accepted sizes of combinations.
+    /// - Note: The range is contracted until empty while iterating over
+    /// combinations of different sizes. When the range is empty, iteration is
+    /// finished.
     @usableFromInline
-    internal var indexes: [Base.Index]
+    internal var kRange: Range<Int>
 
+    /// Whether or not iteration is finished (`kRange` is empty)
     @usableFromInline
-    internal var finished: Bool
+    internal var isFinished: Bool {
+      return kRange.isEmpty
+    }
+    
+    @usableFromInline
+    internal var indexes: [Base.Index]
 
     internal init(_ combinations: Combinations) {
       self.base = combinations.base
-      self.finished = combinations.k < 0
-      self.indexes = combinations.k < 0
-        ? []
-        : Array(combinations.base.indices.prefix(combinations.k))
+      self.kRange = combinations.kRange ?? 0..<0
+      self.indexes = Array(combinations.base.indices.prefix(kRange.lowerBound))
     }
 
     /// Advances the current indices to the next set of combinations. If
@@ -80,24 +123,35 @@ extension Combinations: Sequence {
     ///     // so the iteration is finished.
     @usableFromInline
     internal mutating func advance() {
+      /// Advances `kRange` by incrementing its `lowerBound` until the range is
+      /// empty, when iteration is finished.
+      func advanceKRange() {
+        if kRange.lowerBound < kRange.upperBound {
+          let advancedLowerBound = kRange.lowerBound + 1
+          kRange = advancedLowerBound ..< kRange.upperBound
+          indexes.removeAll(keepingCapacity: true)
+          indexes.append(contentsOf: base.indices.prefix(kRange.lowerBound))
+        }
+      }
+      
       guard !indexes.isEmpty else {
         // Initial state for combinations of 0 elements is an empty array with
         // `finished == false`. Even though no indexes are involved, advancing
         // from that state means we are finished with iterating.
-        finished = true
+        advanceKRange()
         return
       }
-    
+      
       let i = indexes.count - 1
       base.formIndex(after: &indexes[i])
       if indexes[i] != base.endIndex { return }
-
+      
       var j = i
       while indexes[i] == base.endIndex {
         j -= 1
         guard j >= 0 else {
-          // Finished iterating over combinations
-          finished = true
+          // Finished iterating over combinations of this size.
+          advanceKRange()
           return
         }
 
@@ -113,7 +167,7 @@ extension Combinations: Sequence {
 
     @inlinable
     public mutating func next() -> [Base.Element]? {
-      if finished { return nil }
+      guard !isFinished else { return nil }
       defer { advance() }
       return indexes.map { i in base[i] }
     }
@@ -129,10 +183,78 @@ extension Combinations: Equatable where Base: Equatable {}
 extension Combinations: Hashable where Base: Hashable {}
 
 //===----------------------------------------------------------------------===//
-// combinations(count:)
+// combinations(ofCount:)
 //===----------------------------------------------------------------------===//
 
 extension Collection {
+  /// Returns a collection of combinations of this collection's elements, with
+  /// each combination having the specified number of elements.
+  ///
+  /// This example prints the different combinations of 1 and 2 from an array of
+  /// four colors:
+  ///
+  ///     let colors = ["fuchsia", "cyan", "mauve", "magenta"]
+  ///     for combo in colors.combinations(ofCount: 1...2) {
+  ///         print(combo.joined(separator: ", "))
+  ///     }
+  ///     // fuchsia
+  ///     // cyan
+  ///     // mauve
+  ///     // magenta
+  ///     // fuchsia, cyan
+  ///     // fuchsia, mauve
+  ///     // fuchsia, magenta
+  ///     // cyan, mauve
+  ///     // cyan, magenta
+  ///     // mauve, magenta
+  ///
+  /// The returned collection presents combinations in a consistent order, where
+  /// the indices in each combination are in ascending lexicographical order.
+  /// That is, in the example above, the combinations in order are the elements
+  /// at `[0]`, `[1]`, `[2]`, `[3]`, `[0, 1]`, `[0, 2]`, `[0, 3]`, `[1, 2]`,
+  /// `[1, 3]`, and finally `[2, 3]`.
+  ///
+  /// This example prints _all_ the combinations (including an empty array and
+  /// the original collection) from an array of numbers:
+  ///
+  ///     let numbers = [10, 20, 30, 40]
+  ///     for combo in numbers.combinations(ofCount: 0...) {
+  ///         print(combo)
+  ///     }
+  ///     // []
+  ///     // [10]
+  ///     // [20]
+  ///     // [30]
+  ///     // [40]
+  ///     // [10, 20]
+  ///     // [10, 30]
+  ///     // [10, 40]
+  ///     // [20, 30]
+  ///     // [20, 40]
+  ///     // [30, 40]
+  ///     // [10, 20, 30]
+  ///     // [10, 20, 40]
+  ///     // [10, 30, 40]
+  ///     // [20, 30, 40]
+  ///     // [10, 20, 30, 40]
+  ///
+  /// If `kRange` is `0...0`, the resulting sequence has exactly one element, an
+  /// empty array. The given range is limited to `0...base.count`. If the
+  /// limited range is empty, the resulting sequence has no elements.
+  ///
+  /// - Parameter kRange: The range of numbers of elements to include in each
+  /// combination.
+  ///
+  /// - Complexity: O(1) for random-access base collections. O(*n*) where *n*
+  /// is the number of elements in the base collection, since `Combinations`
+  /// accesses the `count` of the base collection.
+  @inlinable
+  public func combinations<R: RangeExpression>(
+    ofCount kRange: R
+  ) -> Combinations<Self> where R.Bound == Int {
+    return Combinations(self, kRange: kRange)
+  }
+  
   /// Returns a collection of combinations of this collection's elements, with
   /// each combination having the specified number of elements.
   ///
@@ -159,7 +281,9 @@ extension Collection {
   ///
   /// - Parameter k: The number of elements to include in each combination.
   ///
-  /// - Complexity: O(1)
+  /// - Complexity: O(1) for random-access base collections. O(*n*) where *n*
+  /// is the number of elements in the base collection, since `Combinations`
+  /// accesses the `count` of the base collection.
   @inlinable
   public func combinations(ofCount k: Int) -> Combinations<Self> {
     assert(k >= 0, "Can't have combinations with a negative number of elements.")