Search code examples
numpymultidimensional-arraytranspose

Transposing a Numpy Array on a slice


I have a 2d array and I am trying to create a 3d array in which each each row is a repeated element, in this case 9 times, of the original array. I think this involves transposing on some kind of a np slice.... my numpy skills are a bit rough...

Here is an example:

input:

an_array = np.array([1,2,3,4,5,6])
a = an_array.reshape(3,2)
a

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

My desired output is as follows:

array([[[1, 1, 1, 1, 1, 1, 1, 1, 1],
        [2, 2, 2, 2, 2, 2, 2, 2, 2]],

        [[3, 3, 3, 3, 3, 3, 3, 3, 3],
        [4, 4, 4, 4, 4, 4, 4, 4, 4]],

        [[5, 5, 5, 5, 5, 5, 5, 5, 5],
        [6, 6, 6, 6, 6, 6, 6, 6, 6]]])

This was my idea, but it does not quite give the desired output. The rows are in the wrong order plus the shape is (2,3,9) instead of (3,2,9), which is an easy issue to resolve, but anyway, I was wondering if there might be quick way to do this?

new = np.transpose([a[:]]*9)
new

array([[[1, 1, 1, 1, 1, 1, 1, 1, 1],
        [3, 3, 3, 3, 3, 3, 3, 3, 3],
        [5, 5, 5, 5, 5, 5, 5, 5, 5]],

       [[2, 2, 2, 2, 2, 2, 2, 2, 2],
        [4, 4, 4, 4, 4, 4, 4, 4, 4],
        [6, 6, 6, 6, 6, 6, 6, 6, 6]]])

Solution

  • reshape and using broadcasting

    a_out = a.reshape(3,-1,1) * np.ones(9)
    
    Out[40]:
    array([[[1., 1., 1., 1., 1., 1., 1., 1., 1.],
            [2., 2., 2., 2., 2., 2., 2., 2., 2.]],
    
           [[3., 3., 3., 3., 3., 3., 3., 3., 3.],
            [4., 4., 4., 4., 4., 4., 4., 4., 4.]],
    
           [[5., 5., 5., 5., 5., 5., 5., 5., 5.],
            [6., 6., 6., 6., 6., 6., 6., 6., 6.]]])
    

    Or

    a_out = np.broadcast_to(np.reshape(a, (3,-1,1)), (3,2,9))
    
    Out[43]:
    array([[[1, 1, 1, 1, 1, 1, 1, 1, 1],
            [2, 2, 2, 2, 2, 2, 2, 2, 2]],
    
           [[3, 3, 3, 3, 3, 3, 3, 3, 3],
            [4, 4, 4, 4, 4, 4, 4, 4, 4]],
    
           [[5, 5, 5, 5, 5, 5, 5, 5, 5],
            [6, 6, 6, 6, 6, 6, 6, 6, 6]]])