Given two integers n and k, return all possible combinations of k numbers chosen from the range [1, n]. You may return the answer in any order.
My code using backtracking is given below.
def combine(n: int, k: int) -> list[list[int]]:
def backtrack(i: int, comb: list[int]) -> None:
if len(comb) == k:
ans.append([*comb])
return
# Number of elements still needed to make a k-combination.
need = k - len(comb)
# The range of numbers is [i, n], therefore, size=n - i + 1
remaining = n - i + 1
# This is the last offset from which we can still make a k-combination.
offset = remaining - need
for j in range(i, i + offset + 1):
comb.append(j)
backtrack(j + 1, comb)
comb.pop()
ans: list[list[int]] = []
backtrack(1, [])
return ans
It works as expected.
Now, I'm looking at LeetCode 1286. Iterator for Combination, which basically asks for an Iterator[list[int]] or Generator[list[int]] instead of returning all the combinations at once.
Technically, LeetCode 1286 works with strings, but for the sake of similarity to LeetCode 77, let's talk about integers only, since it makes no difference to the algorithm.
Can the above code be converted to return an iterator? The reason I'd prefer starting with the above code and not something completely different is because of its simplicity, and to keep the two solutions similar as much as possible.
My research efforts:
I've studied the sample code for itertools.combinations, but it works differently from my code above, and, IMO, is pretty convoluted since it uses some obscure/non-intuitive features such as for-else and loop variables referenced later (Python loop variables remain in scope even after the loop exits).
I've also looked at Algorithm to return all combinations of k elements from n, but due to the highly generic nature of the question (it doesn't specify a programming language or whether the combinations should be returned all at once or as an iterator), the answers are all over the place.
Finally, I've looked at Creating all possible k combinations of n items in C++, which specifies C++, but not an iterator, and thus, doesn't fit my purpose, since I already know how to generate all combinations.
You can pass yielding through recursion with minimal changes to your code.
comb if length is appropriate instead of appending it to ans.backtrack(j + 1, comb) yields after each recursive call.backtrack(1, []) from combine_iterator.def combine_iterator(n: int, k: int) -> list[list[int]]:
def backtrack(i: int, comb: list[int]) -> None:
if len(comb) == k:
yield [*comb]
return
# Number of elements still needed to make a k-combination.
need = k - len(comb)
# The range of numbers is [i, n], therefore, size=n - i + 1
remaining = n - i + 1
# This is the last offset from which we can still make a k-combination.
offset = remaining - need
for j in range(i, i + offset + 1):
comb.append(j)
yield from backtrack(j + 1, comb)
comb.pop()
return backtrack(1, [])
Checking the type of returned object:
type(combine_iterator(5, 3))
>> generator
type(next(combine_iterator(5, 3)))
>> list
type(next(combine_iterator(5, 3))[0])
>> int
Checking that the results are the same for combine and combine_iterator:
for n in range(2, 20):
for k in range(1, n + 1):
res = combine(n, k)
res_it = combine_iterator(n, k)
assert list(res) == list(res_it), f"!!! Different results for n={n}, k={k}"
print("Done!")
Output:
Done!