diff --git a/Guides/Combinations.md b/Guides/Combinations.md index 6f0e186a..b52400f4 100644 --- a/Guides/Combinations.md +++ b/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 diff --git a/README.md b/README.md index 9f07d3d3..40ced02f 100644 --- a/README.md +++ b/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 diff --git a/Sources/Algorithms/Combinations.swift b/Sources/Algorithms/Combinations.swift index 42dde9cb..643feada 100644 --- a/Sources/Algorithms/Combinations.swift +++ b/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 { 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? + + /// 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( + _ 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 { } } - 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 + /// 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( + ofCount kRange: R + ) -> Combinations 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 { assert(k >= 0, "Can't have combinations with a negative number of elements.") diff --git a/Tests/SwiftAlgorithmsTests/CombinationsTests.swift b/Tests/SwiftAlgorithmsTests/CombinationsTests.swift index 4079d6d7..4239b7e0 100644 --- a/Tests/SwiftAlgorithmsTests/CombinationsTests.swift +++ b/Tests/SwiftAlgorithmsTests/CombinationsTests.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 @@ -13,33 +13,114 @@ import XCTest import Algorithms final class CombinationsTests: XCTestCase { + func testCount() { + let c = "ABCD" + + let c0 = c.combinations(ofCount: 0).count + XCTAssertEqual(c0, 1) + + let c1 = c.combinations(ofCount: 1).count + XCTAssertEqual(c1, 4) + + let c2 = c.combinations(ofCount: 2).count + XCTAssertEqual(c2, 6) + + let c3 = c.combinations(ofCount: 3).count + XCTAssertEqual(c3, 4) + + let c4 = c.combinations(ofCount: 4).count + XCTAssertEqual(c4, 1) + + let c5 = c.combinations(ofCount: 0...0).count + XCTAssertEqual(c5, 1) + + let c6 = c.combinations(ofCount: 1...1).count + XCTAssertEqual(c6, 4) + + let c7 = c.combinations(ofCount: 1...2).count + XCTAssertEqual(c7, 10) + + let c8 = c.combinations(ofCount: 1...3).count + XCTAssertEqual(c8, 14) + + let c9 = c.combinations(ofCount: 2...4).count + XCTAssertEqual(c9, 11) + + // `k` greater than element count results in same number of combinations + let c10 = c.combinations(ofCount: 3...10).count + XCTAssertEqual(c10, 5) + + // `k` greater than element count results in same number of combinations + let c11 = c.combinations(ofCount: 4...10).count + XCTAssertEqual(c11, 1) + + // `k` entirely greater than element count results in no combinations + let c12 = c.combinations(ofCount: 5...10).count + XCTAssertEqual(c12, 0) + + let c13 = c.combinations(ofCount: 0...).count + XCTAssertEqual(c13, 16) + + let c14 = c.combinations(ofCount: ...3).count + XCTAssertEqual(c14, 15) + + let c15 = c.combinations(ofCount: 0...).count + XCTAssertEqual(c15, 16) + } + func testCombinations() { let c = "ABCD" - + let c1 = c.combinations(ofCount: 1) XCTAssertEqual(["A", "B", "C", "D"], c1.map { String($0) }) - + let c2 = c.combinations(ofCount: 2) XCTAssertEqual(["AB", "AC", "AD", "BC", "BD", "CD"], c2.map { String($0) }) - + let c3 = c.combinations(ofCount: 3) XCTAssertEqual(["ABC", "ABD", "ACD", "BCD"], c3.map { String($0) }) - + let c4 = c.combinations(ofCount: 4) XCTAssertEqual(["ABCD"], c4.map { String($0) }) + + let c5 = c.combinations(ofCount: 2...4) + XCTAssertEqual(["AB", "AC", "AD", "BC", "BD", "CD", "ABC", "ABD", "ACD", "BCD", "ABCD"], c5.map { String($0) }) + + let c6 = c.combinations(ofCount: 0...4) + XCTAssertEqual(["", "A", "B", "C", "D", "AB", "AC", "AD", "BC", "BD", "CD", "ABC", "ABD", "ACD", "BCD", "ABCD"], c6.map { String($0) }) + + let c7 = c.combinations(ofCount: 0...) + XCTAssertEqual(["", "A", "B", "C", "D", "AB", "AC", "AD", "BC", "BD", "CD", "ABC", "ABD", "ACD", "BCD", "ABCD"], c7.map { String($0) }) + + let c8 = c.combinations(ofCount: ...4) + XCTAssertEqual(["", "A", "B", "C", "D", "AB", "AC", "AD", "BC", "BD", "CD", "ABC", "ABD", "ACD", "BCD", "ABCD"], c8.map { String($0) }) + + let c9 = c.combinations(ofCount: ...3) + XCTAssertEqual(["", "A", "B", "C", "D", "AB", "AC", "AD", "BC", "BD", "CD", "ABC", "ABD", "ACD", "BCD"], c9.map { String($0) }) + + let c10 = c.combinations(ofCount: 1...) + XCTAssertEqual(["A", "B", "C", "D", "AB", "AC", "AD", "BC", "BD", "CD", "ABC", "ABD", "ACD", "BCD", "ABCD"], c10.map { String($0) }) } func testEmpty() { // `k == 0` results in one zero-length combination XCTAssertEqualSequences([[]], "".combinations(ofCount: 0)) + XCTAssertEqualSequences([[]], "".combinations(ofCount: 0...0)) XCTAssertEqualSequences([[]], "ABCD".combinations(ofCount: 0)) - + XCTAssertEqualSequences([[]], "ABCD".combinations(ofCount: 0...0)) + // `k` greater than element count results in zero combinations XCTAssertEqualSequences([], "".combinations(ofCount: 5)) + XCTAssertEqualSequences([], "".combinations(ofCount: 5...10)) XCTAssertEqualSequences([], "ABCD".combinations(ofCount: 5)) + XCTAssertEqualSequences([], "ABCD".combinations(ofCount: 5...10)) } func testCombinationsLazy() { XCTAssertLazySequence("ABC".lazy.combinations(ofCount: 1)) + XCTAssertLazySequence("ABC".lazy.combinations(ofCount: 1...3)) + XCTAssertLazySequence("ABC".lazy.combinations(ofCount: 1...)) + XCTAssertLazySequence("ABC".lazy.combinations(ofCount: ...3)) + XCTAssertLazySequence("ABC".lazy.combinations(ofCount: 0...)) } }