Representing nested loops in Python
You are solving a simple Diophantine equation and for that you use the following Python code.
## 3a+b+c+d=10
r=10/3
for a in range(r, 0, -1):
r=10-3*a
for b in range(r, 0, -1):
r=10-3*a-b
for c in range(r, 0, -1):
d=10-3*a-b-c
if d>0:
print a, b, c, d, 3*a + b + c + d
Keeping the essential nature of the code, how would you represent it "nicely" so that it expands to accommodate more variables in the Diophantine equation?
There are nine solutions:
1 6 1
1 5 2
1 4 3
1 3 4
1 2 5
1 1 6
2 3 1
2 2 2
2 1 3
source to share
I would create a recursive generator function where the arguments are the total s
and multipliers for each element:
def solve(s, multipliers):
if not multipliers:
if s == 0:
yield ()
return
c = multipliers[0]
for i in xrange(s // c, 0, -1):
for solution in solve(s - c * i, multipliers[1:]):
yield (i, ) + solution
for solution in solve(10, [3, 1, 1]):
print solution
Result:
(2, 3, 1) (2, 2, 2) (2, 1, 3) (1, 6, 1) (1, 5, 2) (1, 4, 3) (1, 3, 4) (1, 2, 5) (1, 1, 6)
source to share
You can first determine the possible values of each variable, and then iterate over all the possible combinations using itertool
product
:
from itertools import product
## 3a+b+c+d=10
A = range(10, 0, -1)
B = range(10, 0, -1)
C = range(10, 0, -1)
for a, b, c in product(A, B, C):
d = 10 - 3 * a - b - c
if d > 0:
print a, b, c, d, 3 * a + b + c + d
Output:
2 2 1 1 10 2 1 2 1 10 2 1 1 2 10 1 5 1 1 10 1 4 2 1 10 1 4 1 2 10 1 3 3 1 10 1 3 2 2 10 1 3 1 3 10 1 2 4 1 10 1 2 3 2 10 1 2 2 3 10 1 2 1 4 10 1 1 5 1 10 1 1 4 2 10 1 1 3 3 10 1 1 2 4 10 1 1 1 5 10
Note that by using the same r
for all loops you are doing more work than is really necessary. So it depends on the application if this solution helps.
source to share