k 组合的迭代器

Question

LeetCode 77. 组合:

给定两个整数 n 和 k，返回从范围 [1, n] 中选择的 k 个数字的所有可能组合。 您可以按任意顺序返回答案。

我使用回溯的代码如下。

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

它按预期工作。

现在，我正在查看 LeetCode 1286。组合迭代器，它基本上要求

Iterator[list[int]]

或

Generator[list[int]]

，而不是一次返回所有组合。

从技术上来说，LeetCode 1286 适用于字符串，但为了与 LeetCode 77 相似，我们只讨论整数，因为它对算法没有区别。

上面的代码可以转换为返回一个迭代器吗？我更喜欢从上面的代码开始而不是完全不同的代码，因为它很简单，并且尽可能保持两个解决方案相似。

我的研究成果：

我研究了 itertools.combinations 的示例代码，但它的工作方式与上面的代码不同，而且，IMO，它非常复杂，因为它使用了一些晦涩/不直观的功能，例如

for-else

和引用的循环变量稍后（即使在循环退出后，Python 循环变量仍保留在范围内）。

我还查看了从 n 中返回 k 元素的所有组合的算法，但是由于问题的高度通用性（它没有指定编程语言或是否应该立即返回所有组合或作为迭代器），答案无处不在。

最后，我查看了在 C++ 中创建 n 项的所有可能的 k 组合，它指定了 C++，但不是迭代器，因此不符合我的目的，因为我已经知道如何生成所有组合。

Answer 1

您可以通过递归传递收益，只需对代码进行最少的更改。

如果长度合适，则生成
1. comb

   ，而不是将其附加到 

   ans

   。

产生一切
2. backtrack(j + 1, comb)

   每次递归调用后都会产生。

从
3. backtrack(1, [])

   返回

   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)
            for entity in backtrack(j + 1, comb):
                yield entity
            comb.pop()

    return backtrack(1, [])

检查返回对象的类型：

type(combine_iterator(5, 3))
>> generator

type(next(combine_iterator(5, 3)))
>> list

type(next(combine_iterator(5, 3))[0])
>> int

检查

combine

 和 

combine_iterator

 的结果是否相同：

for n in range(2, 20):
    for k in range(1, n):
        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!")

输出：

Done!