numpy：如何将元素映射到数组

Question

假设我有一个二维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]],
 ...