How to reshape Numpy array with padded 0's

Question

I have a Numpy array that looks like

array([1, 2, 3, 4, 5, 6, 7, 8])

and I want to reshape it to an array

array([[5, 0, 0, 6],
       [0, 1, 2, 0],
       [0, 3, 4, 0],
       [7, 0, 0, 8]])

More specifically, I'm trying to reshape a 2D numpy array to get a 3D Numpy array to go from

array([[ 1,  2,  3,  4,  5,  6,  7,  8],
       [ 9, 10, 11, 12, 13, 14, 15, 16],
       [17, 18, 19, 20, 21, 22, 23, 24],
       ...
       [ 9, 10, 11, 12, 13, 14, 15, 16],
       [89, 90, 91, 92, 93, 94, 95, 96]])

to a numpy array that looks like

array([[[ 5,  0,  0,  6],
        [ 0,  1,  2,  0],
        [ 0,  3,  4,  0],
        [ 7,  0,  0,  8]],

       [[13,  0,  0, 14],
        [ 0,  9, 10,  0],
        [ 0, 11, 12,  0],
        [15,  0,  0, 16]],
       ...
       [[93,  0,  0, 94],
        [ 0, 89, 90,  0],
        [ 0, 91, 92,  0],
        [95,  0,  0, 96]]])

Is there an efficient way to do this using numpy functionality, particularly vectorized?

Answer 1

We can make use of slicing -

def expand(a): # a is 2D array      
    out = np.zeros((len(a),4,4),dtype=a.dtype)
    out[:,1:3,1:3] = a[:,:4].reshape(-1,2,2)
    out[:,::3,::3] = a[:,4:].reshape(-1,2,2)
    return out

The benefit is memory and hence perf. efficiency, as only the output would occupy memory space. The steps involved work with views thanks to the slicing on the input and output.

Sample run -

2D input :

In [223]: a
Out[223]: 
array([[ 1,  2,  3,  4,  5,  6,  7,  8],
       [ 9, 10, 11, 12, 13, 14, 15, 16]])

In [224]: expand(a)
Out[224]: 
array([[[ 5,  0,  0,  6],
        [ 0,  1,  2,  0],
        [ 0,  3,  4,  0],
        [ 7,  0,  0,  8]],

       [[13,  0,  0, 14],
        [ 0,  9, 10,  0],
        [ 0, 11, 12,  0],
        [15,  0,  0, 16]]])

1D input (feed in 2D extended input with None) :

In [225]: a = np.array([1, 2, 3, 4, 5, 6, 7, 8])

In [226]: expand(a[None])
Out[226]: 
array([[[5, 0, 0, 6],
        [0, 1, 2, 0],
        [0, 3, 4, 0],
        [7, 0, 0, 8]]])