How can I make a for-loop pyramid more concise in Python?

PythonPerformancePython 2.7Python 3.x

Python Problem Overview


In solid mechanics, I often use Python and write code that looks like the following:

for i in range(3):
    for j in range(3):
        for k in range(3):
            for l in range(3):
                # do stuff

I do this really often that I start to wonder whether there is a more concise way to do this. The drawback of the current code is: if I comply with PEP8, then I cannot exceed the 79-character-limit per line, and there is not too much space left, especially if this is again in a function of a class.

Python Solutions


Solution 1 - Python

By using nested for-loops you're basically trying to create what's known as the (Cartesian) product of the input iterables, which is what the product function, from itertools module, is for.

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

And in your code you can do :

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

Based on python documentation, "This function is roughly equivalent to the following code, except that the actual implementation does not build up intermediate results in memory:"

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)

Solution 2 - Python

The idea to use itertools.product is a good one. Here's a more general approach that will support ranges of varying sizes.

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

Solution 3 - Python

It won't be more concise as it will cost you a generator function, but at least you won't be bothered by 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

(in python 2.x you should use xrange instead of range in the generator function)

EDIT:

Above method should be fine when the depth of the pyramid is known. But you can also make a generic generator that way without any external module :

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

(I used (l, k, j, i) because in above generator, first index varies first)

Solution 4 - Python

This is equivalent:

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)

If you're doing this all the time, consider using a general generator for it:

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)

Attributions

All content for this solution is sourced from the original question on Stackoverflow.

The content on this page is licensed under the Attribution-ShareAlike 4.0 International (CC BY-SA 4.0) license.

Content TypeOriginal AuthorOriginal Content on Stackoverflow
QuestionYuxiang WangView Question on Stackoverflow
Solution 1 - PythonMazdakView Answer on Stackoverflow
Solution 2 - PythonFMcView Answer on Stackoverflow
Solution 3 - PythonSerge BallestaView Answer on Stackoverflow
Solution 4 - PythonL3viathanView Answer on Stackoverflow