用2D ndarray给出的值填充3D ndarray中平方矩阵的对角线

Question

给定z具有形状的3D ndarray (k,n,n)，是否可以不使用迭代就v用具有形状的2D ndarray给出的值填充k
nxn矩阵的对角线(k,n)？

例如，运算结果应与循环遍历k个矩阵相同：

z = np.zeros((3,10,10))
v = np.arange(30).reshape((3,10))

for i in range(len(z)): 
    np.fill_diagonal(z[i], v[i])

有没有一种方法可以避免np.fill_diagonal在循环内反复调用？如果可能的话，我宁愿能够适用于更高维的数组，以及，一个解决方案，其中z.shape == (a,b,c,...,k,n,n)和v.shape = (a,b,c,...,k,n)

Answer 1

这是通用的n-dim数组-

diag_view = np.einsum('...ii->...i',z)
diag_view[:] = v

另一个重塑-

n = v.shape[-1] 
z.reshape(-1,n**2)[:,::n+1] = v.reshape(-1,n)
# or z.reshape(z.shape[:-2]+(-1,))[...,::n+1] = v

另一个masking-

m = np.eye(n, dtype=bool) # n = v.shape[-1] from earlier
z[...,m] = v

初始化输出z

如果我们需要初始化输出数组z，并且要覆盖一般的n-dim情况，则应该是：

z = np.zeros(v.shape + (v.shape[-1],), dtype=v.dtype)

然后，我们继续前面列出的方法。