numpy：如何将（N，）维向量插入（N，M，D）维数组作为沿D轴的新元素？ （N，M，D + 1）

Question

a is an array of shape (N, M, D) == (20, 4096, 6) .

b is an array of shape (N,) == (20,) .

I would like to insert the values of b to a such that each value in b is appended element-wise to the D dim in a (7th element in a ).

So c would be such an array, of shape (20, 4096, 7) , where c[i,:,-1] == b[i] for all i , and c[...,:-1] == a .

I know you could just make a new array and add the values accordingly eg:

N, M, D = a.shape # (20, 4096, 6)
c = np.zeros((N, M, D+1))
c[...,:-1] = a
for i in range(N):
    c[i,:,-1] = b[i]

But was wondering if one of the numpy wizards here had a more slick way of doing this with numpy ops and no intermediate arrays.

Answer 1

Replicate b along the second axis after extending it to 3D and then concatenate with a along the last axis -

b_rep = np.repeat(b[:,None,None],a.shape[1],axis=1)
out = np.concatenate((a, b_rep,axis=-1)

Alternatively, we can use np.broadcast_to to create the replicated version :

b_rep = np.broadcast_to(b[:,None,None], (len(b), a.shape[1],1)

Answer 2

Here is another one-liner

np.r_['2,3,0', a, np.broadcast_to(b, (a.T.shape[1:])).T]

Also, I'd like to mention that your original method is actually close to the (or at least a) recommended way. Just use empty instead of zeros and broadcasting instead of the loop:

res = np.empty((N,M,D+1), np.promote_types(a.dtype, b.dtype))
res[..., :-1], res[..., -1] = a, b[:, None]

...

And - just for fun - one more, which I expressly do not recommend. Do not use this!

np.where(np.arange(D+1)<D, np.lib.stride_tricks.as_strided(a, (N,M,D+1), a.strides), b[:, None, None])