Skip to content
Advertisement

numpy filling the diagonal 0 for 3D array

Suppose I have a 3D array, how can I fill the diag of the first two dimensions to zero. For example

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

Is there a way to replace the for loop?

Advertisement

Answer

The following is one of the solution

a = np.random.rand(2,2,3)
np.einsum('iij->ij',a)[...] = 0
User contributions licensed under: CC BY-SA
8 People found this is helpful
Advertisement