Numpy: как вставить (N,) -мерный вектор в (N, M, D) -мерный массив как новые элементы вдоль оси D? (N, M, D+1)

Question

a это массив формы (N, M, D) == (20, 4096, 6),

b это массив формы (N,) == (20,),

Я хотел бы вставить значения b в a так что каждое значение в b добавляется поэлементно к D тусклый a (7-й элемент в a).

Так c был бы такой массив, формы (20, 4096, 7), где c[i,:,-1] == b[i] для всех i, а также c[...,:-1] == a,

Я знаю, что вы могли бы просто создать новый массив и добавить значения соответственно, например:

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]

Но мне было интересно, есть ли у кого-нибудь из этих волшебников-нудистов более плавный способ сделать это с помощью numpy ops и без промежуточных массивов.

Answer 1

Копировать b вдоль второй оси после расширения его до 3D а затем объединить с a вдоль последней оси -

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

В качестве альтернативы мы можем использовать np.broadcast_to создать реплицированную версию:

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

Answer 2

Вот еще один вкладыш

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

Кроме того, я хотел бы отметить, что ваш оригинальный метод на самом деле близок (или, по крайней мере, a) рекомендован. Просто используйте empty вместо zeros и вещание вместо цикла:

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

...

И - просто для удовольствия - еще один, который я категорически не рекомендую. Не используйте это!

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