How can I make a for-loop pyramid more concise in Python?
Question:
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.
Answers:
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)
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)
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)
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
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.
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)
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)
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)
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