python - 使用 Itertools 的等效嵌套循环结构

Question

Python 的简洁语法通过它的电池允许用可读的一行来表达冗长的代码行。考虑以下示例

====================================================|
for a in range(3):                                  |
    for b in range(3):                              |
        for c in range(3):                          |
            print (a,b,c),                          |
-  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -|
for e in product(range(3), repeat=3):               |
    print e,                                        |
====================================================|
for a in range(3):                                  |
    for b in range(a , 3):                          |
        for c in range(b , 3):                      |
            print (a,b,c),                          |
-  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -|
for e in combinations_with_replacement(range(3), 3):|
    print e,                                        |
====================================================|
for a in range(3):                                  |
    for b in range(a + 1, 3):                       |
        for c in range(b + 1, 3):                   |
            print (a,b,c),                          |
-  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -|
for e in combinations(range(3), 3):                 |
    print e,                                        |
====================================================|
for a in range(3):                                  |
    for b in range(3):                              |
        for c in range(3):                          |
            if len(set([a,b,c])) == 3:              |
                print (a,b,c),                      |
-  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -|
for e in permutations(range(3)):                    |
    print e,                                        |
====================================================|

最后我得到了一个深层嵌套的依赖循环，我试图简洁地表达但失败了

循环结构如下

for a in A():
    for b in B(a):
        for c in C(b):
            foo(a,b,c)

这样的结构可以用等效的 itertools 符号表示吗？

Answer 1

没有确切的 itertools 解决方案，但 itertools 函数的简单组合就足够了:

def chain_imap_accumulate(seq, f):
    def acc_f(x):
        for n in f(x[-1]):
            yield x + (n,)
    return chain.from_iterable(imap(acc_f, seq))

def accumulative_product(*generators):
    head, tail = generators[0], generators[1:]
    head = imap(tuple, head())
    return reduce(chain_imap_accumulate, tail, head)

快速测试。定义:

from itertools import chain, imap, izip
chain_ = chain.from_iterable

def A():
    yield 'A'
    yield 'B'

def B(x):
    yield int(x, 16)
    yield int(x, 16) + 1

def C(x):
    yield str(x) + 'Z'
    yield str(x) + 'Y'

结果:

>>> list(accumulative_product(A, B, C))
[('A', 10, '10Z'), ('A', 10, '10Y'), 
 ('A', 11, '11Z'), ('A', 11, '11Y'), 
 ('B', 11, '11Z'), ('B', 11, '11Y'), 
 ('B', 12, '12Z'), ('B', 12, '12Y')]

几乎所有的复杂性都来自于输入的积累，如上述代码的快速“推导”所示。 final (c) 值可以使用几个嵌套的 itertools 结构生成:

>>> list(chain_(imap(C, chain_(imap(B, (A()))))))
['10Z', '10Y', '11Z', '11Y', '11Z', '11Y', '12Z', '12Y']

这可以用reduce 概括。要使用 reduce，chain_imap 不能使用标准的 imap 参数顺序。它必须被交换:

def chain_imap(seq, f):
    return chain.from_iterable(imap(f, seq))

这给出了相同的结果:

>>> list(reduce(chain_imap, [B, C], A()))
['10Z', '10Y', '11Z', '11Y', '11Z', '11Y', '12Z', '12Y']

最后的任务是累积初始值，以便您可以访问 a、b 和 c。这需要一些思考才能正确，但实现相当简单——我们只需将 f 转换为一个函数，该函数忽略除最后一个输入值之外的所有输入值，并将新值附加到完整值输入:

def chain_imap_accumulate(seq, f):
    def acc_f(x):
        for n in f(x[-1]):
            yield x + (n,)
    return chain.from_iterable(imap(acc_f, seq))

这需要将第一个输入包装在元组中，因此我们将 A 映射到 tuple:

>>> list(reduce(chain_imap_accumulate, [B, C], imap(tuple, A())))
[('A', 10, '10Z'), ('A', 10, '10Y'), 
 ('A', 11, '11Z'), ('A', 11, '11Y'), 
 ('B', 11, '11Z'), ('B', 11, '11Y'), 
 ('B', 12, '12Z'), ('B', 12, '12Y')]

为清楚起见，重写上面的代码，结果是此答案顶部的代码。

顺便说一句，chain_imap_accumulate 可以使用 genex 更简洁地重写。这可以与 accumulative_product 的较短版本结合使用，以获得非常紧凑的定义(如果您对这类事情感兴趣)。这也恰好完全消除了 itertools 依赖:

def chain_map_accumulate(seq, f):
    return (x + (n,) for x in seq for n in f(x[-1]))

def accumulative_product2(*gens):
    return reduce(chain_map_accumulate, gens[1:], (tuple(x) for x in gens[0]()))