Generating a nested list with x dimensions

user_14396787 0
Viewed 56

I want to create a function that returns a matrix with a given shape, similar to what np.zeros(shape) would do in numpy.

I tried the following, which works:

shape = [1,2,3]
x = [ lambda : 0 ]
x.append( lambda : [ x[0]() for _ in range(shape[0])] )
x.append( lambda : [ x[1]() for _ in range(shape[1])] )
x.append( lambda : [ x[2]() for _ in range(shape[2])] )
print(x[-1]())

But when I try to automize this to work with any amount of dimensions as follows, I get a RecursionError.

shape = [1,2,3]
x = [ lambda : 0 ]
for i,s in enumerate(shape): 
    x.append( lambda: [ x[i]() for _ in range(s)] )
print(x[-1]())

Does someone know what went wrong here, or if there is a better way to do this?

python
1 Answers

The problem with the second version of your code is that, when the lambdas get executed, they will use the current values of i and s, which are the ones they got during the last loop (here, 2 and 3).

So, the ones you created in the loop will actually all be

lambda : [ x[2]() for _ in range(3)]

including the one at x[2], which will call itself in an infinite recursion.

The classic 'trick' is to take advantage of the fact that default arguments are evaluated at definition time. So, we use default arguments for the lambdas, which will 'capture' the value of i and s when each is defined:

shape = [1, 2, 3]
x = [lambda: 0]
for i,s in enumerate(shape): 
    x.append(lambda i=i, s=s: [ x[i]() for _ in range(s)])

print(x[-1]())

Output:

[[[0], [0]], [[0], [0]], [[0], [0]]]

You could also build it using a recursive function, which might be easier to understand:

def nested_list(shape, value):
    if len(shape) == 1:
        return [value] *  shape[0]
    return [nested_list(shape[:-1], value) for _ in range(shape[-1])]


shape = [1,2,3]
print(nested_list(shape, 0))
# [[[0], [0]], [[0], [0]], [[0], [0]]]
user_5500940
Related