Numpy array of integers into numpy array of arrays

Question

I would like to convert numpy array into numpy array of arrays.

I have an array: a = [[0,0,0],[0,255,0],[0,255,255],[255,255,255]]

and I would like to have: b = [[[0,0,0],[0,0,0],[0,0,0]],[[0,0,0],[255,255,255],[0,0,0]],[[0,0,0],[255,255,255],[255,255,255]],[[255,255,255],[255,255,255],[255,255,255]]]

Is there any easy way to do it?

I have tried with np.where(a == 0, [0,0,0],[255,255,255]) but I got the following error:

ValueError: operands could not be broadcast together with shapes

Answer 1

You can use broadcast_to as

b = np.broadcast_to(a, (3,4,3))

where a was shape (3,4). Then you need to swap the axes around

import numpy as np
a = np.array([[0,0,0],[0,255,0],[0,255,255],[255,255,255]])
b = np.broadcast_to(a, (3,4,3))
c = np.moveaxis(b, [0,1,2], [2,0,1])
c

giving

array([[[  0,   0,   0],
        [  0,   0,   0],
        [  0,   0,   0]],

       [[  0,   0,   0],
        [255, 255, 255],
        [  0,   0,   0]],

       [[  0,   0,   0],
        [255, 255, 255],
        [255, 255, 255]],

       [[255, 255, 255],
        [255, 255, 255],
        [255, 255, 255]]])

A more direct method broadcasting method suggested by @Divakar is

 b = np.broadcast(a[:,:,None], (4,3,3))

which produces the same output without axis swapping.