问题描述
假设我有一个 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]]])