使用 2D ndarray 给出的值填充 3D ndarray 内方阵的对角线 [英] Filling the diagonals of square matrices inside a 3D ndarray with values given by a 2D ndarray

问题描述

给定一个形状为 (k,n,n) 的 3D ndarray z，是否可以不使用迭代来填充具有给定值的 k nxn 矩阵的对角线由形状为 (k,n)?

Given a 3D ndarray z with shape (k,n,n), is it possbile without using iteration to fill the diagonals of the k nxn matrices with values given by a 2D ndarray v with shape (k,n)?

例如，操作的结果应该与循环 k 个矩阵相同:

For example, the result of the operation should be the same as looping over k matrices:

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)

Is there an way to do this without repeatedly calling np.fill_diagonal inside a loop? If possible, I would prefer a solution that can be applied to arrays of higher dimensions as well, where z.shape == (a,b,c,...,k,n,n) and v.shape = (a,b,c,...,k,n)

推荐答案

这里是通用的 n-dim 数组 -

Here's for generic n-dim arrays -

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 情况的数组，它将是:

If we need to initialize the output array z and one that will cover for generic n-dim cases, it would be :

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

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

Then, we proceed with the earlier listed approaches.

这篇关于使用 2D ndarray 给出的值填充 3D ndarray 内方阵的对角线的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！