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

Posted 2023-02-25

【中文标题】使用 Itertools 具有内部中断的可变级别的嵌套循环

【英文标题】：Variable levels of nested loops with inner breaks using Itertools

【发布时间】：2017-02-13 06:04:04

【问题描述】：

经验丰富的开发人员学习 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 循环，使用递归或迭代。在编写代码之前，我搜索了一下，发现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() 以便它可以执行并可以选择跳出最里面的循环。

任何 Pythonic 建议？我讨厌必须显式地编写变量嵌套循环或遍历 product() 的返回并手动检查子集。这看起来很老派:-)

【问题讨论】：

不要做for subset in list(product(*sets)):...省略list，这会迫使你将整个product物化到内存中，并避免延迟迭代product的内存效率。由于您没有保留它，它只是毫无意义的低效。

如果您使用 product 或 combinations 之类的 itertools，它们总是会生成所有组合。你不能让它做一个内部break。 您可以通过存储“失败”案例、检查最后一个元素并使用continue 移动到下一个来跳过组合，但您仍然会得到每个元素。

我仍然不明白使用product 有什么帮助。

@juanpa.arrivillaga 是的，谢谢，这是有道理的（关于不使用 list()）。我也看到我不能做我想做的事。每个解决方案都涉及生成所有子集并自己跟踪它们。在这种情况下，显式编码嵌套循环至少可以节省所有这些迭代。

可能可以编写一个自定义生成器，该生成器可以由send-ing 在值中发出信号，但使用起来可能有点尴尬。

【参考方案1】：

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

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 实际上会提前退出内部循环，如果你告诉它，它会完全跳过不必要的迭代。

【参考方案2】：

这是一个相当基本的迭代有序组合生成器，它采用回调函数。它不像 BrenBarn 的解决方案那样通用，但它确实跳过了问题中指定的生成产品：如果在传递 arg seq 时回调失败，返回 False（或错误的东西），那么 combo 将跳过生成那些以seq[:-1]开头的其他子序列。

def combo(n, k, callback):
    a = list(range(k))
    ok = callback(a)

    k -= 1
    while True:
        i = k
        if not ok: i -= 1
        while True:
            a[i] += 1
            if a[i] == (n + i - k):
                i -= 1
                if i < 0: 
                    return
            else:
                break 
        v = a[i] + 1  
        a[i+1:] = range(v, v + k - i) 
        ok = callback(a)

# test

n = 8
k = 4

def do_something(seq):
    s = sum(seq)
    ok = s <= seq[-2] * 3
    print(seq, s, ok)
    return ok

combo(n, k, do_something)

输出

[0, 1, 2, 3] 6 True
[0, 1, 2, 4] 7 False
[0, 1, 3, 4] 8 True
[0, 1, 3, 5] 9 True
[0, 1, 3, 6] 10 False
[0, 1, 4, 5] 10 True
[0, 1, 4, 6] 11 True
[0, 1, 4, 7] 12 True
[0, 1, 5, 6] 12 True
[0, 1, 5, 7] 13 True
[0, 1, 6, 7] 14 True
[0, 2, 3, 4] 9 True
[0, 2, 3, 5] 10 False
[0, 2, 4, 5] 11 True
[0, 2, 4, 6] 12 True
[0, 2, 4, 7] 13 False
[0, 2, 5, 6] 13 True
[0, 2, 5, 7] 14 True
[0, 2, 6, 7] 15 True
[0, 3, 4, 5] 12 True
[0, 3, 4, 6] 13 False
[0, 3, 5, 6] 14 True
[0, 3, 5, 7] 15 True
[0, 3, 6, 7] 16 True
[0, 4, 5, 6] 15 True
[0, 4, 5, 7] 16 False
[0, 4, 6, 7] 17 True
[0, 5, 6, 7] 18 True
[1, 2, 3, 4] 10 False
[1, 2, 4, 5] 12 True
[1, 2, 4, 6] 13 False
[1, 2, 5, 6] 14 True
[1, 2, 5, 7] 15 True
[1, 2, 6, 7] 16 True
[1, 3, 4, 5] 13 False
[1, 3, 5, 6] 15 True
[1, 3, 5, 7] 16 False
[1, 3, 6, 7] 17 True
[1, 4, 5, 6] 16 False
[1, 4, 6, 7] 18 True
[1, 5, 6, 7] 19 False
[2, 3, 4, 5] 14 False
[2, 3, 5, 6] 16 False
[2, 3, 6, 7] 18 True
[2, 4, 5, 6] 17 False
[2, 4, 6, 7] 19 False
[2, 5, 6, 7] 20 False
[3, 4, 5, 6] 18 False
[3, 4, 6, 7] 20 False
[3, 5, 6, 7] 21 False
[4, 5, 6, 7] 22 False

如果您注释掉 if not ok: i -= 1 行，它将产生所有组合。

---

可以轻松修改此代码以进行更大的跳过。我们没有使用回调的布尔返回，而是返回想要的跳过级别。如果它返回零，那么我们不做任何跳过。如果它返回 1，那么我们跳过以seq[:-1] 开头的子序列，如上面的版本。如果回调返回 2，那么我们跳过以seq[:-2] 等开头的子序列。

def combo(n, k, callback):
    a = list(range(k))
    skips = callback(a)

    k -= 1
    while True:
        i = k - skips
        while True:
            a[i] += 1
            if a[i] == (n + i - k):
                i -= 1
                if i < 0: 
                    return
            else:
                break 
        v = a[i] + 1  
        a[i+1:] = range(v, v + k - i) 
        skips = callback(a)