python - 使用 Itertools 实现具有内部中断的可变级别嵌套循环

Question

经验丰富的开发人员学习Python。

我一次从大小为 n 的列表中迭代组合 k。我一直在使用

from itertools import chain, combinations
for subset in (combinations(range(n), k)) : 
    doSomethingWith(subset)

现在的问题是，大多数时候我的 doSomethingWith() 效率不高，我想尽可能多地跳过它们。如果 doSomthingWith() 对于给定的子集失败，我可以跳过最右边坐标较大的每个子集。换句话说，如果 (1,3,5,8) 失败，那么我想查看的下一个子集是 (1,3,6,0)，跳过 (1,3,5,9), (1, 3,5,10)等

我意识到我需要对循环索引进行显式控制。我需要使用递归或迭代的可变数量的嵌套 for 循环。在编码之前，我用 Google 搜索了一下，发现 this thread这看起来很有希望。

所以现在我有:

from itertools import product, repeat
set  = range(n)
sets = repeat(set, k)

for subset in list(product(*sets)) :
    doSomethingWith(subset)

漂亮的Pythonic，但我仍然有同样的问题。我无法告诉 Product() 何时跳出最内层循环。我真正想要的是能够将回调函数传递到 Product() 中，以便它可以执行并可选择跳出最内层循环。

有什么 Python 风格的建议吗？我不想显式地编写变量嵌套循环或迭代product()的返回值并手动检查子集。这看起来很老派:-)

Answer 1

这是一个有趣的问题。我炮制了一个生成器，可以用来实现你想要的东西。这更像是一个原型(prototype)，而不是完整的解决方案。它可能需要进行调整才能真正有用，并且可能存在我没有考虑到的边缘情况。 (一方面，现在它无法在可能耗尽的任意可迭代对象上正常工作，只能在像列表这样的可重复迭代对象上正常工作。)

class SkipUp(Exception):
    def __init__(self, numSkips):
        self.numSkips = numSkips
        super(SkipUp, self).__init__(numSkips)

def breakableProduct(*sets):
    if not sets:
        yield []
        return
    first, rest = sets[0], sets[1:]
    for item in first:
        subProd = breakableProduct(*rest)
        for items in subProd:
            try:
                yield [item] + items
            except SkipUp as e:
                if e.numSkips == 0:
                    yield None
                    break
                else:
                    e.numSkips -= 1
                    yield subProd.throw(e)

您可以或多或少像普通的itertools.product一样使用breakableProduct:

>>> prod = breakableProduct([1, 2, 3], [11, 22, 33], [111, 222, 333])
... for x, y, z in prod:
...     print(x, y, z)
1 11 111
1 11 222
1 11 333
1 22 111
1 22 222
# ...etc...
3 33 222
3 33 333

但是，在从中获取值后，如果您愿意，可以使用 .throw 传入 SkipUp 异常，该异常的参数是要跳到其下一个元素的集合的索引。也就是说，如果您想跳过第三个集合的所有元素并跳转到第二个集合的下一个元素，请使用 SkipUp(1) (1 是第二个集合，因为索引是从 0 开始的):

>>> prod = breakableProduct([1, 2, 3], [11, 22, 33], [111, 222, 333])
... for x, y, z in prod:
...     print(x, y, z)
...     if z == 222:
...         prod.throw(SkipUp(1))
1 11 111
1 11 222
1 22 111
1 22 222
1 33 111
1 33 222
2 11 111
2 11 222
2 22 111
2 22 222
2 33 111
2 33 222
3 11 111
3 11 222
3 22 111
3 22 222
3 33 111
3 33 222

看看这如何中止 222 之后的最内层迭代，而不是向前跳转到中间生成器的下一个迭代。如果你想一直跳到最外层迭代:

>>> prod = breakableProduct([1, 2, 3], [11, 22, 33], [111, 222, 333])
... for x, y, z in prod:
...     print(x, y, z)
...     if z == 222:
...         prod.throw(SkipUp(0))
1 11 111
1 11 222
2 11 111
2 11 222
3 11 111
3 11 222

因此 SkipUp(0) 跳转到第一个迭代器的下一个“tick”(即列表 [1, 2, 3])。加入 SkipUp(2) 不会产生任何效果，因为它仅意味着“跳到最内层迭代器的下一次迭代”，而这正是常规迭代会执行的操作。

与使用来自 itertools 的 product 或 combinations 之类的解决方案相比，此解决方案的优势在于您无法阻止这些生成器生成每一个组合。因此，即使你想跳过某些元素，itertools 仍然会执行所有循环来生成所有你不想要的元素，并且只有在它们生成后你才会丢弃它们。另一方面，如果您告诉它，这个 breakableProduct 实际上会提前退出内部循环，因此它将完全跳过不必要的迭代。