How to replace each element of an array with another array in Python?

MWE: Let us consider the following example.

L0=[[b,0],[b,b]], L1=[[b,b],[b,1]], L2=[[b,b],[2,b]]

S=[[0,1,2],[2,0,1]]

Is there any way the replace each element of S by L0 for 0 and L1 for 1 and L2 for 2 in S to get S1 like in the image? enter image description here

Actually, I want a python program which will check: if the element of S is zero then it will replace 0 by the predefined 2D array and so on.

Solution

Yes. We can first construct a numpy array that contains L0, L1 and L2:

A = np.array([L0, L1, L2])

Next we construct a numpy array of S:

B = np.array(S)

now we have for C = A[B] (or C = np.take(A,B,axis=0) as suggested by @Divakar):

>>> C = np.take(A,B,axis=0)
>>> C
array([[[[b, 0],
         [b, b]],

        [[b, b],
         [b, 1]],

        [[b, b],
         [2, b]]],


       [[[b, b],
         [2, b]],

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

        [[b, b],
         [b, 1]]]])

This is of course not exactly what we intended: we want to obtain a 2D-array. We can do this by first transposing (or swapaxes, like @PaulPanzer suggests) and then reshaping, we obtain:

>>> C.transpose(0,2,1,3).reshape(4,6)
array([[b, 0, b, b, b, b],
       [b, b, b, 1, 2, b],
       [b, b, b, 0, b, b],
       [2, b, b, b, b, 1]])

Since 4 and 6 of course depend on the size of the dimensions of L0, L1, L2 and S, we can also calculate them based on that size:

A = np.array([L0, L1, L2])
B = np.array(S)
m, n = B.shape
_, u, v = A.shape
np.take(A,B,axis=0).swapaxes(1,2).reshape(u*m, v*n)

Like @DSM says, from Numpy-1.13, there is np.block function for this purpose, and we can write it as:

>>> np.block([[A[i] for i in row] for row in S])
array([[b, 0, b, b, b, b],
       [b, b, b, 1, 2, b],
       [b, b, b, 0, b, b],
       [2, b, b, b, b, 1]])