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]]])