steeeve steeeve - 1 month ago 11
Python Question

Return multiple matrices based on user input to function

I have a function that accepts user parameter k and generates a matrix of the form:

a = np.array((k,2,5))


Leaving the function I would like to return k 2x5 matrices. Essentially,

if k == 2:
return a[0,:,:], a[1,:,:]
if k == 3:
return a[0,:,:], a[1,:,:], a[2,:,:]


Is there a way to do this dynamically without having to hard-code each of the possible values of k?

Thanks.

Answer

Sure:

return tuple(a[i,:,:] for i in range(k))

Python's great! 8-))