How to get indices of non-diagonal elements of a numpy array?
To get the mask, you can use np.eye
, like so -
~np.eye(a.shape[0],dtype=bool)
To get the indices, add np.where
-
np.where(~np.eye(a.shape[0],dtype=bool))
Sample run -
In [142]: a
Out[142]:
array([[7412, 33, 2],
[ 2, 7304, 83],
[ 3, 101, 7237]])
In [143]: ~np.eye(a.shape[0],dtype=bool)
Out[143]:
array([[False, True, True],
[ True, False, True],
[ True, True, False]], dtype=bool)
In [144]: np.where(~np.eye(a.shape[0],dtype=bool))
Out[144]: (array([0, 0, 1, 1, 2, 2]), array([1, 2, 0, 2, 0, 1]))
There are few more ways to get such a mask for a generic non-square input array.
With np.fill_diagonal
-
out = np.ones(a.shape,dtype=bool)
np.fill_diagonal(out,0)
With broadcasting
-
m,n = a.shape
out = np.arange(m)[:,None] != np.arange(n)
>>> import numpy as np
>>> a = np.array([[7412, 33, 2],
... [2, 7304, 83],
... [3, 101, 7237]])
>>> non_diag = np.ones(shape=a.shape, dtype=bool) - np.identity(len(a)).astype(bool)
>>> non_diag
array([[False, True, True],
[ True, False, True],
[ True, True, False]], dtype=bool)