Search code examples
pythonnumpynumpy-ndarrayarray-broadcasting

How to create a 3-d array from 1-d array?


Let's say I have a 1d array

a = np.array([1, 2, 3])

What's the best way to get the array b with shape (3, 4, 5) from a? Every value of the array a is used to initialize a 4x5 array and stacking all these arrays will create the array b. I was wondering if I could avoid looping to create the array b.

b = np.array([[[1, 1, 1, 1, 1], [1, 1, 1, 1, 1], [1, 1, 1, 1, 1], [1, 1, 1, 1, 1]], [[2, 2, 2, 2, 2], [2, 2, 2, 2, 2], [2, 2, 2, 2, 2], [2, 2, 2, 2, 2]], [[3, 3, 3, 3, 3], [3, 3, 3, 3, 3], [3, 3, 3, 3, 3], [3, 3, 3, 3, 3]]])

Output for b:

array([[[1, 1, 1, 1, 1],
    [1, 1, 1, 1, 1],
    [1, 1, 1, 1, 1],
    [1, 1, 1, 1, 1]],
   [[2, 2, 2, 2, 2],
    [2, 2, 2, 2, 2],
    [2, 2, 2, 2, 2],
    [2, 2, 2, 2, 2]],
   [[3, 3, 3, 3, 3],
    [3, 3, 3, 3, 3],
    [3, 3, 3, 3, 3],
    [3, 3, 3, 3, 3]]])

Solution

  • You can achieve by creating an array of ones and multiplying it by a[:,None,None] to use broadcasting.

    import numpy as np
    
    a = np.array([1, 2, 3])
    b = np.ones((3,4,5), dtype=int)*a[:,None,None]
    

    You can also do this without creating an array of ones using np.broadcast_to.

    import numpy as np
    
    a = np.array([1, 2, 3])
    b = np.broadcast_to(a[:,None,None], (3,4,5))
    

    There should also be a way to do this using np.repeat. Currently, I only found a way that achieves the desired result using nested calls.

    b = np.repeat(np.repeat(a[:,None], 5, axis=1)[:,None], 4, axis=1)