如何在Python中简化循环金字塔？

Question

在实体力学中，我经常使用Python并编写类似于以下内容的代码：for i in range(3):    for j in range(3):        for k in range(3):            for l in range(3):                # do stuff我经常这样做，以至于我开始怀疑是否有更简洁的方法可以做到这一点。当前代码的缺点是：如果我遵守PEP8，那么我就不能超过每行79个字符的限制，并且没有太多的空间，特别是如果这又是一个类的函数。

Answer 1

根据您要执行的操作，您可以使用该itertools模块来最大程度地减少for循环（或zip）。在这种情况下，itertools.product将创建您使用4个循环完成的工作：

>>> list(product(range(3),repeat=4))

[(0, 0, 0, 0), (0, 0, 0, 1), (0, 0, 0, 2), (0, 0, 1, 0), (0, 0, 1, 1),

 (0, 0, 1, 2), (0, 0, 2, 0), (0, 0, 2, 1), (0, 0, 2, 2), (0, 1, 0, 0),

 (0, 1, 0, 1), (0, 1, 0, 2), (0, 1, 1, 0), (0, 1, 1, 1), (0, 1, 1, 2),

 (0, 1, 2, 0), (0, 1, 2, 1), (0, 1, 2, 2), (0, 2, 0, 0), (0, 2, 0, 1),

 (0, 2, 0, 2), (0, 2, 1, 0), (0, 2, 1, 1), (0, 2, 1, 2), (0, 2, 2, 0),

 (0, 2, 2, 1), (0, 2, 2, 2), (1, 0, 0, 0), (1, 0, 0, 1), (1, 0, 0, 2),

 (1, 0, 1, 0), (1, 0, 1, 1), (1, 0, 1, 2), (1, 0, 2, 0), (1, 0, 2, 1),

 (1, 0, 2, 2), (1, 1, 0, 0), (1, 1, 0, 1), (1, 1, 0, 2), (1, 1, 1, 0),

 (1, 1, 1, 1), (1, 1, 1, 2), (1, 1, 2, 0), (1, 1, 2, 1), (1, 1, 2, 2),

 (1, 2, 0, 0), (1, 2, 0, 1), (1, 2, 0, 2), (1, 2, 1, 0), (1, 2, 1, 1),

 (1, 2, 1, 2), (1, 2, 2, 0), (1, 2, 2, 1), (1, 2, 2, 2), (2, 0, 0, 0),

 (2, 0, 0, 1), (2, 0, 0, 2), (2, 0, 1, 0), (2, 0, 1, 1), (2, 0, 1, 2),

 (2, 0, 2, 0), (2, 0, 2, 1), (2, 0, 2, 2), (2, 1, 0, 0), (2, 1, 0, 1),

 (2, 1, 0, 2), (2, 1, 1, 0), (2, 1, 1, 1), (2, 1, 1, 2), (2, 1, 2, 0),

 (2, 1, 2, 1), (2, 1, 2, 2), (2, 2, 0, 0), (2, 2, 0, 1), (2, 2, 0, 2),

 (2, 2, 1, 0), (2, 2, 1, 1), (2, 2, 1, 2), (2, 2, 2, 0), (2, 2, 2, 1),

 (2, 2, 2, 2)]

在您的代码中，您可以执行以下操作：

for i,j,k,l in product(range(3),repeat=4):

    #do stuff

此函数等效于以下代码，不同之处在于实际实现不会在内存中建立中间结果：

def product(*args, **kwds):

    # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy

    # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111

    pools = map(tuple, args) * kwds.get('repeat', 1)

    result = [[]]

    for pool in pools:

        result = [x+[y] for x in result for y in pool]

    for prod in result:

        yield tuple(prod)

编辑：正如@ PeterE在评论中所说，product()即使范围的长度不同，也可以使用：

product(range(3),range(4),['a','b','c'] ,some_other_iterable)

Answer 2

使用的想法itertools.product是一个好主意。这是一种更通用的方法，将支持各种大小的范围。

from itertools import product

def product_of_ranges(*ns):

    for t in product(*map(range, ns)):

        yield t

for i, j, k in product_of_ranges(4, 2, 3):

    # do stuff

Answer 3

它不会更加简洁，因为它将花费您一个生成器函数，但是至少您不会被PEP8所困扰：

def tup4(n):

    for i in range(n):

        for j in range(n):

            for k in range(n):

                for l in range(n):

                    yield (i, j, k, l)

for (i, j, k, l) in tup4(3):

    # do your stuff

（在python 2.x中，您应该使用xrange而不是range在generator函数中）

编辑：

当金字塔的深度已知时，上述方法应该很好。但是您也可以通过这种方式生成通用生成器，而无需任何外部模块：

def tup(n, m):

    """ Generate all different tuples of size n consisting of integers < m """

    l = [ 0 for i in range(n)]

    def step(i):

        if i == n : raise StopIteration()

        l[i] += 1

        if l[i] == m:

            l[i] = 0

            step(i+ 1)

    while True:

        yield tuple(l)

        step(0)

for (l, k, j, i) in tup(4, 3):

    # do your stuff

（我用过，(l, k, j, i)因为在上面的生成器中，第一个索引首先变化）