如何在Python中使for循环金字塔更简洁？

Question

如何在Python中使for循环金字塔更简洁？

62

在固体力学中，我经常使用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个字符的限制，而且左右的空间不多，尤其是如果这又是类的函数中的代码。

- Yuxiang Wang

你只是在迭代范围吗？那么有一种更短的方法（虽然不一定更易读）。 - L3viathan

5

如果一个算法的时间复杂度为O(n^4)，那么它就是O(n^4)，没有任何变通的余地。如果想避免79个字符的限制，可以考虑将代码分割成多个函数。这样做可以显著提高代码的可读性和可测试性。 - UltraInstinct

6

嗯，深度嵌套循环不是编程的好方式...所以我认为你应该更关注避免深度嵌套循环而不是担心PEP8。 - sarveshseri

4

使用向量化的numpy操作，比如numpy.einsum()，参考使用NumPy实现快速张量旋转。 - jfs

重复？这绝对看起来是更好的问题... - Zizouz212

4个回答

15

使用 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

- FMc

13

虽然这样会让你需要使用生成器函数，但是它不会更为简洁，但至少你不必被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)。

编辑：

当金字塔的深度已知时，上述方法应该没问题。但是你也可以创建一个通用的生成器，而无需任何外部模块：

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)，因为在上面的生成器中，第一个索引首先变化。

- Serge Ballesta

8

这是等效的：

for c in range(3**4):
    i = c // 3**3 % 3
    j = c // 3**2 % 3
    k = c // 3**1 % 3
    l = c // 3**0 % 3
    print(i,j,k,l)

如果你经常这样做，考虑使用一个通用的生成器:

def nestedLoop(n, l):
    return ((tuple((c//l**x%l for x in range(n-1,-1,-1)))) for c in range(l**n))

for (a,b,c,d) in nestedLoop(4,3):
    print(a,b,c,d)

- L3viathan

1

在这种情况下，为了说明问题，您可能还应该使用 c // 3**3 % 3、c // 3**2 % 3 等等。 ;-) - martineau

1

...这就是生成器。 - L3viathan

请编辑文本后再进行翻译，谢谢。 - L3viathan

1

同意...但这似乎对大多数人来说并不像简洁那么重要。您可以通过yield一个元组而不是返回生成器表达式来进一步缩短它，同时通过将reversed(range(n))更改为range(n)[::-1]（或range(n,-1,-1)）来进一步缩短它。 - martineau

1

@martineau 后者，因为切片操作会生成一个列表。 - L3viathan

显示剩余6条评论

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Mazdak · Accepted Answer

使用嵌套的for循环，您基本上是在尝试创建所谓的输入可迭代对象的（笛卡尔）积，这就是来自itertools模块的product函数的作用。

>>> 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),
...

而在您的代码中，您可以这样做：

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

根据Python文档，“该函数大致等同于以下代码，但实际实现不会在内存中构建中间结果：”

def product(*args, repeat=1):
    # 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 = [tuple(pool) for pool in args] * repeat
    result = [[]]
    for pool in pools:
        result = [x+[y] for x in result for y in pool]
    for prod in result:
        yield tuple(prod)