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

Question

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

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

有没有办法替换for循环？

标签： pythonnumpynumpy-ndarray

Answer 1

以下是解决方案之一

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