numpy 填充 3D 数组的对角线 0

如何解决numpy 填充 3D 数组的对角线 0

假设我有一个 3D 数组，如何将前两个维度的 diag 填充为零。例如

a = np.random.rand(2,2,3)
for i in range(3):
    np.fill_diagonal(a[:,:,i],0)

有没有办法替换 for 循环？

解决方法

以下是解决方案之一

a = np.random.rand(2,2,3)
np.einsum('iij->ij',a)[...] = 0

,

np.diag 函数返回一个二维对角矩阵。

a[:,:,0] = np.diag((1,1))

,

In [6]: a = np.random.randint(1,10,(2,3))
   ...: for i in range(3):
   ...:     np.fill_diagonal(a[:,i],0)
In [7]: a
Out[7]: 
array([[[0,0],[7,4,4]],[[8,7],[0,0]]])

索引对角线很容易 - 只需为两个维度使用相同的索引数组。因此我们刚刚设置的 0 是：

In [8]: idx=np.arange(2)
In [9]: a[idx,idx,:]
Out[9]: 
array([[0,0]])

并用于设置一个值：

In [10]: a[idx,:] = 10
In [11]: a
Out[11]: 
array([[[10,10],[ 7,[[ 8,[10,10]]])