0

I want to create a function that behaves the same way as the for loops down there, dependant on the variable n:

function(n, max_anz, max_size_ , step)

n = 1

z = 0
for a1 in range(max_anz + 1):
    for x1 in range(0, max_size, step):
        
        print([a1, x1])
        z += 1
print(z)

n = 2

z = 0
for a1 in range(max_anz + 1):
    for x1 in range(0, max_size, step):
        
        for a2 in range(max_anz + 1):
            for x2 in range(0, max_size, step):
                
                print([a1, x1, a2, x2])
                z = z+1
print(z)

n = 3

z = 0
for a1 in range(max_anz + 1):
    for x1 in range(0, max_size, step):
        
        for a2 in range(max_anz + 1):
            for x2 in range(0, max_size, step):
                
                for a3 in range(max_anz + 1):
                    for x3 in range(0, max_size, step):
                
                        print([a1, x1, a2, x2, a3, x3])
                        z = z+1

print(z)

Ousbond
  • 11
  • 4
  • Do you really want to print the items and their count in the function, or get the items in a list? – hpchavaz Nov 27 '21 at 11:13
  • Similar questions: [Function with varying number of For Loops (python)](https://stackoverflow.com/questions/7186518/function-with-varying-number-of-for-loops-python); [How to make n-dimensional nested for-loops in Python?](https://stackoverflow.com/questions/50860578/how-to-make-n-dimensional-nested-for-loops-in-python); [Dynamically generating a flexible number of nested for loops](https://stackoverflow.com/questions/52456836/dynamically-generating-a-flexible-number-of-nested-for-loops); – Stef Nov 27 '21 at 16:56
  • Similar questions: [Variable number of predictable for loops in Python](https://stackoverflow.com/questions/32059848/variable-number-of-predictable-for-loops-in-python); [Arbitrary number of nested loops dependent on the previous loop in Python](https://stackoverflow.com/questions/35659647/arbitrary-number-of-nested-loops-dependent-on-the-previous-loop-in-python); [Variable length nested loop Python](https://stackoverflow.com/questions/49404539/variable-length-nested-loop-python) – Stef Nov 27 '21 at 16:57

3 Answers3

3

Using itertools.product with its 'repeat' argument.

import itertools as it

def function(n, max_anz, max_size, step):
    lst = [list(e) for e in it.product(range(max_anz+1), range(0, max_size, step), repeat = n)]
    print(*lst, sep='\n')
    print(len(lst))

Or with a generator to prevent building the full list:

def function(n, max_anz, max_size, step):
    g = it.product(range(max_anz+1), range(0, max_size, step), repeat = n)
    
    i = 0
    for e in g:
        i += 1
        print(list(e))
    print(i)

which could be shortened:

def function(n, max_anz, max_size, step):
    i = 0
    for e it.product(range(max_anz+1), range(0, max_size, step), repeat = n):
        i += 1
        print(list(e))
    print(i)
hpchavaz
  • 1,368
  • 10
  • 16
  • I recommend adding `i = -1` before the for-loop. Otherwise, you might be trying to print i+1 without having defined i at all. For instance, `function(1, 1, 0, 1)` currently raises Exception `UnboundLocalError: local variable 'i' referenced before assignment`. In general, I have an extremely strong distate for referencing the loop variables outside of the loop. If you want to count the number of iterations, you can do that explicitly: `number_of_iterations=0; for e in it.product(...): print(list(e)); number_of_iterations += 1; print(number_of_iterations)` – Stef Nov 27 '21 at 16:14
  • 1
    @Stef, I edited the answer following your suggestion. Feel free to edit. – hpchavaz Nov 27 '21 at 16:26
1

You could do what you want using recursion, with a function that calls itself with n-1 (if n is greater than one) and passes in some partial results to build on.

The snippet below demonstrates this approach on a simpler problem with only a single for loop at each level of n.

def make_list(n, limit, context=None):
    z = 0
    for val in range(limit):
        new_list = [] if context is None \
            else context.copy()
        new_list.append(val)
        if n > 1:
            z += make_list(n-1, limit, context=new_list)
        else:
            print(new_list)
            z += 1
    return z

make_list(2, 3)
# [0, 0]
# [0, 1]
# [0, 2]
# [1, 0]
# [1, 1]
# [1, 2]
# [2, 0]
# [2, 1]
# [2, 2]
# Out[10]: 9
TMBailey
  • 557
  • 3
  • 14
1

itertools.product is the go-to solution for this in python, but as a programmer, you shouldn't be sitting around waiting for features to arrive in your chosen language or library. For example, if you were to take this problem to JavaScript, how would you do it?

You notice a pattern, so wrap it in a function and define the necessary parameters. Using simple inductive reasoning we can structure the function -

  1. If n is less than or equal to zero, no more loops, yield the empty result
  2. (inductive) n is 1 or greater. Run your loops, and for each result r of the subproblem n-1, prepend a and x with the result
def myfunc(n, max_anz, max_size, step):
  if n <= 0:
    yield ()                                   #1
  else:
    for a in range(max_anz + 1):               #2
      for x in range(0, max_size, step):
        for r in myfunc(n - 1, max_anz, max_size, step):
          yield (a, x, *r)

Now you can call myfunc with your arguments -

for r in myfunc(2, 3, 4, 1):
  print(r)
(0, 0, 0, 0)
(0, 0, 0, 1)
(0, 0, 0, 2)
...
(3, 3, 3, 1)
(3, 3, 3, 2)
(3, 3, 3, 3)

You can gather all of the combinations in a list if you like, where len gives you the total count of combinations (z in your question) -

result = list(myfunc(2, 3, 4, 1))
print(result)
print(len(result))
[
  (0, 0, 0, 0),
  (0, 0, 0, 1),
  (0, 0, 0, 2),
  ...,
  (3, 3, 3, 1),
  (3, 3, 3, 2),
  (3, 3, 3, 3)
]
256
Mulan
  • 129,518
  • 31
  • 228
  • 259