Add range-based permutations(ofCount:)

Author: mdznr

Sources/Algorithms/Permutations.swift

@@ -14,20 +14,56 @@ public struct Permutations<Base: Collection> {
   /// The base collection to iterate over for permutations.
   public let base: Base
 
+  @usableFromInline
   internal let baseCount: Int
-  internal let countToChoose: Int
 
+  /// The range of accepted sizes of permutations.
+  /// - Note: This may be empty if the attempted range entirely exceeds the
+  /// bounds of the size of the `base` collection.
+  @usableFromInline
+  internal let kRange: Range<Int>
+  
+  /// Initializes a `Permutations` for all permutations of `base` of size `k`.
+  /// - Parameters:
+  ///   - base: The collection to iterate over for permutations
+  ///   - k: The expected size of each permutation, or `nil` (default) to
+  ///   iterate over all permutations of the same size as the base collection.
+  @usableFromInline
   internal init(_ base: Base, k: Int? = nil) {
+    let kRange: ClosedRange<Int>?
+    if let countToChoose = k {
+      kRange = countToChoose ... countToChoose
+    } else {
+      kRange = nil
+    }
+    self.init(base, kRange: kRange)
+  }
+  
+  /// Initializes a `Permutations` for all combinations of `base` of sizes
+  /// within a given range.
+  /// - Parameters:
+  ///   - base: The collection to iterate over for permutations.
+  ///   - kRange: The range of accepted sizes of permutations, or `nil` to
+  ///   iterate over all permutations of the same size as the base collection.
+  @usableFromInline
+  internal init<R: RangeExpression>(
+    _ base: Base, kRange: R?
+  ) where R.Bound == Int {
     self.base = base
     let baseCount = base.count
     self.baseCount = baseCount
-    self.countToChoose = k ?? baseCount
+    let upperBound = baseCount + 1
+    self.kRange = kRange?.relative(to: 0 ..< .max)
+      .clamped(to: 0 ..< upperBound) ??
+      baseCount ..< upperBound
   }
 
+  /// The total number of permutations.
+  @inlinable
   public var count: Int {
-    return baseCount >= countToChoose
-      ? stride(from: baseCount, to: baseCount - countToChoose, by: -1).reduce(1, *)
-      : 0
+    return kRange.map {
+      stride(from: baseCount, to: baseCount - $0, by: -1).reduce(1, *)
+    }.reduce(0, +)
   }
 }
 
@@ -38,36 +74,30 @@ extension Permutations: Sequence {
     internal var base: Base
 
     @usableFromInline
-    internal var indexes: [Base.Index]
-    @usableFromInline
-    internal var hasMorePermutations: Bool
+    internal let baseCount: Int
+    
+    /// The current range of accepted sizes of permutations.
+    /// - Note: The range is contracted until empty while iterating over
+    /// permutations of different sizes. When the range is empty, iteration is
+    /// finished.
     @usableFromInline
-    internal var countToChoose: Int = 0
-        
-    /// `true` if we're generating permutations of the full collection.
+    internal var kRange: Range<Int>
+    
+    /// Whether or not iteration is finished (`kRange` is empty)
     @usableFromInline
-    internal var permutesFullCollection: Bool {
-      countToChoose == indexes.count
+    internal var isFinished: Bool {
+      return kRange.isEmpty
     }
 
     @usableFromInline
-    internal init(_ base: Base) {
-      self.base = base
-      self.indexes = Array(base.indices)
-      self.countToChoose = self.indexes.count
-      self.hasMorePermutations = true
-    }
+    internal var indexes: [Base.Index]
 
     @usableFromInline
-    internal init(_ base: Base, count: Int) {
-      self.base = base
-      self.countToChoose = count
-      
-      // Produce exactly one empty permutation when `count == 0`.
-      self.indexes = count == 0 ? [] : Array(base.indices)
-
-      // Can't produce any permutations when `count > base.count`.
-      self.hasMorePermutations = count <= indexes.count
+    internal init(_ permutations: Permutations) {
+      self.base = permutations.base
+      self.baseCount = permutations.baseCount
+      self.kRange = permutations.kRange
+      self.indexes = Array(permutations.base.indices)
     }
 
     /// Advances the `indexes` array such that the first `countToChoose`
@@ -83,6 +113,7 @@ extension Permutations: Sequence {
     /// - Complexity: O(*n*), where *n* is the length of the collection.
     @usableFromInline
     internal mutating func nextState() -> Bool {
+      let countToChoose = self.kRange.lowerBound
       let edge = countToChoose - 1
 
       // Find first index greater than the one at `edge`.
@@ -113,31 +144,53 @@ extension Permutations: Sequence {
 
     @inlinable
     public mutating func next() -> [Base.Element]? {
-      if !hasMorePermutations { return nil }
+      guard !isFinished else { return nil }
+      
+      /// 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 = Array(base.indices)
+        }
+      }
+      
+      let countToChoose = self.kRange.lowerBound
+      if countToChoose == 0 {
+        defer {
+          advanceKRange()
+        }
+        return []
+      }
 
+      let permutesFullCollection = (countToChoose == baseCount)
       if permutesFullCollection {
         // If we're permuting the full collection, each iteration is just a
         // call to `nextPermutation` on `indexes`.
-        defer { hasMorePermutations = indexes.nextPermutation() }
+        defer {
+          let hasMorePermutations = indexes.nextPermutation()
+          if !hasMorePermutations {
+            advanceKRange()
+          }
+        }
         return indexes.map { base[$0] }
       } else {
         // Otherwise, return the items at the first `countToChoose` indices and
         // advance the state.
-        defer { hasMorePermutations = nextState() }
+        defer {
+          let hasMorePermutations = nextState()
+          if !hasMorePermutations {
+            advanceKRange()
+          }
+        }
         return indexes.prefix(countToChoose).map { base[$0] }
       }
     }
   }
 
-  @usableFromInline
-  internal var permutesFullCollection: Bool {
-    baseCount == countToChoose
-  }
-
   public func makeIterator() -> Iterator {
-    permutesFullCollection
-      ? Iterator(base)
-      : Iterator(base, count: countToChoose)
+    Iterator(self)
   }
 }
 
@@ -196,6 +249,13 @@ where Self: BidirectionalCollection, Element: Comparable
 //===----------------------------------------------------------------------===//
 
 extension Collection {
+  @inlinable
+  public func permutations<R: RangeExpression>(
+    ofCount kRange: R
+  ) -> Permutations<Self> where R.Bound == Int {
+    return Permutations(self, kRange: kRange)
+  }
+  
   /// Returns a collection of the permutations of this collection of the
   /// specified length.
   ///
@@ -242,6 +302,7 @@ extension Collection {
   ///   of this entire collection.
   ///
   /// - Complexity: O(1)
+  @inlinable
   public func permutations(ofCount k: Int? = nil) -> Permutations<Self> {
     precondition(
       k ?? 0 >= 0,