Добавление столбца в более высокое измерение

Question

давайте предположим, что у меня есть массив с именем A

A = np.zeros((4, 3, 2))

array([[[0., 0.],
    [0., 0.],
    [0., 0.]],
   [[0., 0.],
    [0., 0.],
    [0., 0.]],
   [[0., 0.],
    [0., 0.],
    [0., 0.]],
   [[0., 0.],
    [0., 0.],
    [0., 0.]]])

и другой массив с именем B

B = np.arange(4)
array([0, 1, 2, 3])

, и я хочу сделать что-то вроде конкатенации в третьем измерении, чтобы получить этот результат :

array([[[0., 0., 0.0],
    [0., 0., 0.0],
    [0., 0., 0.0]],
   [[0., 0., 1.0],
    [0., 0., 1.0],
    [0., 0., 1.0]],
   [[0., 0., 2.0],
    [0., 0., 2.0],
    [0., 0., 2.0]],
   [[0., 0., 3.0],
    [0., 0., 3.0],
    [0., 0., 3.0]]])

я пробовал сервальные способы сделать это, но у меня ничего не получилось.
кто я могу сделать это хорошим способом, не зацикливаясь?

Answer 1

Чтобы добавить дополнительное измерение, вы можете использовать np.append. Вы просто должны получить правильную форму. Вы можете использовать np.repeat() для создания повторяющихся элементов:

A = np.zeros((4, 3, 2))
h, w, d = A.shape
B = np.repeat(np.arange(h), w).reshape([h, w, 1])

np.append(A, B, axis=2)

Вывод:

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

       [[0., 0., 1.],
        [0., 0., 1.],
        [0., 0., 1.]],

       [[0., 0., 2.],
        [0., 0., 2.],
        [0., 0., 2.]],

       [[0., 0., 3.],
        [0., 0., 3.],
        [0., 0., 3.]]])

Answer 2

Мы могли бы передать B соответствующей форме и использовать расширенное индексирование здесь и назначить B, транслируемое по соответствующим осям:

np.concatenate([A, np.broadcast_to(B[:,None,None], A[...,-1:].shape)], -1)

---

print(A)

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

       [[0., 0., 1.],
        [0., 0., 1.],
        [0., 0., 1.]],

       [[0., 0., 2.],
        [0., 0., 2.],
        [0., 0., 2.]],

       [[0., 0., 3.],
        [0., 0., 3.],
        [0., 0., 3.]]])