使用Itertools

时间:2017-02-13 06:04:04

标签: python nested-loops itertools

经验丰富的开发人员学习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(),以便它可以执行并可选地中断最内层循环。

任何Pythonic建议?我不想必须显式编码变量嵌套循环或迭代从product()返回并手动检查子集。这看起来很老了:-)

2 个答案:

答案 0 :(得分: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)

您可以使用breakableProduct或多或少的正常itertools.product

>>> 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)将没有任何效果,因为它只是意味着跳到内部迭代器的下一次迭代",这是常规迭代无论如何都要做的。

使用来自product的{​​{1}}或combinations之类的解决方案的优点是,您无法阻止这些生成器生成每个组合。因此,即使你想跳过一些元素,itertools仍然会进行所有循环以生成你不想要的所有元素,并且你只会在它们已经生成之后丢弃它们。另一方面,这个itertools实际上会提前退出内部循环,因此它将完全跳过不必要的迭代。

答案 1 :(得分:1)

这是一个相当基本的迭代有序组合制作器,它具有回调函数。它不像BrenBarn的解决方案那样通用,但是它确实跳过了问题中指定的产生:如果在传递arg seq时回调失败,返回False(或者假的是什么),那么{{ 1}}将跳过生成以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 行,它将生成所有组合。

可以轻松修改此代码以执行更大的跳过。我们不是使用回调中的布尔返回值来返回所需的跳过级别。如果它返回零,那么我们不做任何跳过。如果它返回1,那么我们跳过以if not ok: i -= 1开头的子序列,如上面的版本所示。如果回调返回2,那么我们跳过以seq[:-1]等开始的子序列

seq[:-2]