python - 使用 numpy 重复数组维度(不使用 np.repeat)

Question

我想复制一个 numpy 数组维度，但要保证原始维度数组和复制维度数组的总和仍然相同。例如考虑 n x m形状数组( a )，我想将其转换为 n x n x m ( b ) 数组，因此 a[i,j] == b[i,i,j] 。不幸的是np.repeat和np.resize不适合这份工作。是否有另一个我可以使用的 numpy 函数，或者这可以通过一些创意索引来实现吗？

>>> import numpy as np
>>> a = np.asarray([1, 2, 3])
>>> a
array([1, 2, 3])
>>> a.shape
(3,)
# This is not what I want...
>>> np.resize(a, (3, 3))
array([[1, 2, 3],
       [1, 2, 3],
       [1, 2, 3]])

在上面的例子中，我想得到这样的结果:

array([[1, 0, 0],
       [0, 2, 0],
       [0, 0, 3]])

Answer 1

从一维数组到二维数组，您可以使用np.diagflat方法，该方法创建一个二维数组，并将展平的输入作为对角线:

import numpy as np
a = np.asarray([1, 2, 3])

np.diagflat(a)
#array([[1, 0, 0],
#       [0, 2, 0],
#       [0, 0, 3]])

更一般地说，您可以创建一个零数组并使用高级索引就地分配值:

a = np.asarray([[1, 2, 3], [4, 5, 6]])

result = np.zeros((a.shape[0],) + a.shape)
idx = np.arange(a.shape[0])
result[idx, idx, :] = a

result
#array([[[ 1.,  2.,  3.],
#        [ 0.,  0.,  0.]],

#       [[ 0.,  0.,  0.],
#        [ 4.,  5.,  6.]]])