list comprehension variable assignment

Question

I am attempting to create a 4d array and assign variables to each of the cells.

Typically I would use four "for loops" but this is very messy and takes up a lot of space.

What i'm currently doing:

for x in range(2):
    for y in range(2):
        for j in range(2):
            for k in range(2):
                array[x,y,j,k] = 1 #will be a function in reality

I've tried using list comprehension but this only creates the list and does not assign variables to each cell.

Are there space-efficient ways to run through multiple for loops and assign variables with only a few lines of code?

Answer 1

Assuming you've already created an empty (numpy?) array, you can use itertools.product to fill it with values:

import itertools

for x, y, j, k in itertools.product(range(2), repeat=4):
    arr[x,y,j,k] = 1

If not all of the array's dimensions are equal, you can list them individually:

for x, y, j, k in itertools.product(range(2), range(2), range(2), range(2)):
    arr[x,y,j,k] = 1