Numpy中沿axis=0重复数组元素的复杂变换问题及解决方案

问题说明 #

给定数组：

arr = np.array([[1,2,3], [4,5,6]])

需要将axis=0上的元素重复2次，期望输出如下：

array([[[1., 4.],
    [1., 4.]],

   [[2., 5.],
    [2., 5.]],

   [[3., 6.],
    [3., 6.]]])

尝试的代码及错误输出 #

使用以下代码尝试实现：

np.ones((2,2))*arr.T[:,:,None]

得到的输出不符合预期：

array([[[1., 1.],
    [4., 4.]],

   [[2., 2.],
    [5., 5.]],

   [[3., 3.],
    [6., 6.]]])

修正后的解决方案 #

通过调整数组的轴顺序即可解决问题，最终代码如下：

np.transpose(np.ones((2,2))*arr.T[:,:,None], axes = tuple([0, 2, 1]))

内容的提问来源于stack exchange，提问作者thesecond