Numpy: как отобразить элемент в массив

Question

Предположим, у меня есть массив 2D numpy, скажем 5*3.Теперь я хотел бы сопоставить каждый элемент i в нем с новым массивом [i, i*i], поэтому полученный массив будет 5*3*2.

. Какой самый эффективный (и элегантный) способ достижения этой цели?

Наивное решение с использованием for:

a = np.arange(15).reshape(5, 3)
r = []
for row in a:
   _row = []
   for i in row:
      _row.append([i, i*i])
   r.append(_row)
return np.array(r)

Answer 1

Вы можете использовать np.dstack для укладки обоих массивов по глубине:

np.dstack([a, a**2])

---

a = np.arange(15).reshape(5, 3)

array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11],
       [12, 13, 14]])

np.dstack([a, a**2])

array([[[  0,   0],
        [  1,   1],
        [  2,   4]],

       [[  3,   9],
        [  4,  16],
        [  5,  25]],
 ...