I have a numpy array of shape (5, 4, 3)
and another numpy array of shape (4,)
and what I want to do is expand the last dimension of the first array
(5, 4, 3) -> (5, 4, 4)
and then broadcast the other array with shape (4,)
such that it fills up the new array cells respectively.
Example:
np.ones((5,4,3)) array([[[1., 1., 1.], [1., 1., 1.], [1., 1., 1.], [1., 1., 1.]], [[1., 1., 1.], [1., 1., 1.], [1., 1., 1.], [1., 1., 1.]], [[1., 1., 1.], [1., 1., 1.], [1., 1., 1.], [1., 1., 1.]], [[1., 1., 1.], [1., 1., 1.], [1., 1., 1.], [1., 1., 1.]], [[1., 1., 1.], [1., 1., 1.], [1., 1., 1.], [1., 1., 1.]]])
becomes
array([[[1., 1., 1., 0.], [1., 1., 1., 0.], [1., 1., 1., 0.], [1., 1., 1., 0.]], [[1., 1., 1., 0.], [1., 1., 1., 0.], [1., 1., 1., 0.], [1., 1., 1., 0.]], [[1., 1., 1., 0.], [1., 1., 1., 0.], [1., 1., 1., 0.], [1., 1., 1., 0.]], [[1., 1., 1., 0.], [1., 1., 1., 0.], [1., 1., 1., 0.], [1., 1., 1., 0.]], [[1., 1., 1., 0.], [1., 1., 1., 0.], [1., 1., 1., 0.], [1., 1., 1., 0.]]])
And then I have another array
array([2., 3., 4., 5.])
which I somehow broadcast with the first one to fill the zeros:
array([[[1., 1., 1., 2.], [1., 1., 1., 3.], [1., 1., 1., 4.], [1., 1., 1., 5.]], [[1., 1., 1., 2.], [1., 1., 1., 3.], [1., 1., 1., 4.], [1., 1., 1., 5.]], [[1., 1., 1., 2.], [1., 1., 1., 3.], [1., 1., 1., 4.], [1., 1., 1., 5.]], [[1., 1., 1., 2.], [1., 1., 1., 3.], [1., 1., 1., 4.], [1., 1., 1., 5.]], [[1., 1., 1., 2.], [1., 1., 1., 3.], [1., 1., 1., 4.], [1., 1., 1., 5.]]])
How can I accomplish this?
Advertisement
Answer
You can use numpy.c_
and numpy.tile
:
A = np.ones((5,4,3), dtype='int') B = np.array([2, 3, 4, 5]) np.c_[A, np.tile(B[:,None], (A.shape[0], 1, 1))]
output:
array([[[1, 1, 1, 2], [1, 1, 1, 3], [1, 1, 1, 4], [1, 1, 1, 5]], [[1, 1, 1, 2], [1, 1, 1, 3], [1, 1, 1, 4], [1, 1, 1, 5]], [[1, 1, 1, 2], [1, 1, 1, 3], [1, 1, 1, 4], [1, 1, 1, 5]], [[1, 1, 1, 2], [1, 1, 1, 3], [1, 1, 1, 4], [1, 1, 1, 5]], [[1, 1, 1, 2], [1, 1, 1, 3], [1, 1, 1, 4], [1, 1, 1, 5]]])
How it works:
# reshape B to add one dimension >>> B[:, None] array([[2], [3], [4], [5]]) # tile to match A's first dimension >>> np.tile(B[:,None], (A.shape[0], 1, 1)) array([[[2], [3], [4], [5]], [[2], [3], [4], [5]], [[2], [3], [4], [5]], [[2], [3], [4], [5]], [[2], [3], [4], [5]]])