python-用1D数组用零填充制作3D数组，具体取决于1D数组的索引

Question

Consider the following 1D arrays

a=np.arange(3)+9
b=np.arange(3)+5

currently I am initializing the new 3d array by using

n=4
cols=3
k=np.vstack((a,b,a*b,np.zeros((n,cols)),a,b,a,a,b**2,np.zeros((n,cols)),a*2,a)).T.reshape(-1,2,n+5)

where a and b will always be the same shape

which results in

array([[[  9.,   5.,  45.,   0.,   0.,   0.,   0.,   9.,   5.],
        [  9.,   9.,  25.,   0.,   0.,   0.,   0.,  18.,   9.]],

       [[ 10.,   6.,  60.,   0.,   0.,   0.,   0.,  10.,   6.],
        [ 10.,  10.,  36.,   0.,   0.,   0.,   0.,  20.,  10.]],

       [[ 11.,   7.,  77.,   0.,   0.,   0.,   0.,  11.,   7.],
        [ 11.,  11.,  49.,   0.,   0.,   0.,   0.,  22.,  11.]]])

How would i use a similar technique, also without a for loop, to change the zero padding to the following:

array([[[  9.,   5.,  45.,   9.,   5.,   0.,   0.,   0.,   0.],
        [  9.,   9.,  25.,  18.,   9.,   0.,   0.,   0.,   0.]],

       [[ 10.,   6.,  60.,   0.,   0.,  10.,   6.,   0.,   0.],
        [ 10.,  10.,  36.,   0.,   0.,  20.,  10.,   0.,   0.]],

       [[ 11.,   7.,  77.,   0.,   0.,   0.,   0.,  11.,   7.],
        [ 11.,  11.,  49.,   0.,   0.,   0.,   0.,  22.,  11.]]])

Answer 1

One can use advanced-indexing to assign those array values into a zeros initialized array given the column indices -

out = np.zeros((3,2,9),dtype=bool)
vals = np.array([[a,b,a*b,a,b],[a,a,b**2,2*a,a]])
out[np.arange(3)[:,None],:, idx] = vals.T

Sample run -

In [448]: a
Out[448]: array([ 9, 10, 11])

In [449]: b
Out[449]: array([5, 6, 7])

In [450]: out
Out[450]: 
array([[[  9.,   5.,  45.,   9.,   5.,   0.,   0.,   0.,   0.],
        [  9.,   9.,  25.,  18.,   9.,   0.,   0.,   0.,   0.]],

       [[ 10.,   6.,  60.,   0.,   0.,  10.,   6.,   0.,   0.],
        [ 10.,  10.,  36.,   0.,   0.,  20.,  10.,   0.,   0.]],

       [[ 11.,   7.,  77.,   0.,   0.,   0.,   0.,  11.,   7.],
        [ 11.,  11.,  49.,   0.,   0.,   0.,   0.,  22.,  11.]]])