How to create a function with a variable number of 'for' loops, each with a distinct index?

63 views Asked by At

The Problem:

Consider a d-dimensional simple cubic lattice.

If the lattice has width L, then the number of lattice sites is Ld. I want to create a list that contains all the positions of the lattice sites, for a general d and L.

For example, when L = 2 and d = 2 this would be [(0, 0), (1, 0), (0, 1), (1, 1)].

My Attempt:

Whilst I can do this for general L, I have not been able to generalise the dimension d.

Below is my solution for d = 3 using three for loops.

def Positions(L):
    PositionList = []
    for i in range(L):
        for j in range(L):
            for k in range(L):
                PositionList.append([k, j, i])
    return PositionList

It's easy to see how I would change it to increase or decrease the dimension d as I could simply add or remove the for loops, but obviously it is very tedious to write this out for large d.

I thought about using a recursive for loop, so I would only be using one for loop for any d, but I don't know how to do this whilst preserving the indexing I need to write out the positions.

In Conclusion:

Is it possible to have a variable number of for loops, each one having a distinct index to solve this problem?

Or is there a better solution without using for loops?

3

There are 3 answers

1
Seyi Daniel On BEST ANSWER

One easy way is using itertools cartesian product:

from itertools import product 
L, D = 2, 2 
print(list(product(list(range(L)), repeat = D)))

Result

[(0, 0), (0, 1), (1, 0), (1, 1)]
4
trincot On

Recursion is indeed the way to go

The idea is:

If you assume your function works for d-1 dimensions, then you can take that result and append to each of the results the value of i (the loop variable), and do that repeatedly for each value of i.

The base case is when d=0, in that case you have just a single, empty result.

Here is how that can be coded:

def Positions(L, d):
    if d == 0:  # base case
        return [[]]
    return [
        [i] + res  # prepend i to the results coming from recursion
            for i in range(L)
                for res in Positions(L, d-1)
    ]

If you are not familiar with the list-comprehension syntax used in the final statement, then here is how you would do it without that syntax:

def Positions(L, d):
    if d == 0:  # base case
        return [[]]
    positions = []
    for i in range(L):
        for res in Positions(L, d-1):
            positions.append([i] + res)
    return positions
0
tiger li On

you use recursion. the first part is the base case and the second part is to add every number from 0 to L-1 for every term in the lattice for the lower dimension

def positions(L,d):
  if d==0:
    return [()]
  else:
    return [(x,)+positions(L,d-1)[y] for x in range(L) for y in range(len(positions(L,d-1)))]