CODEWITHSUNDEEP
pythonnumpymultidimensional-array
drew_is_good
drew_is_good

Reputation: 163

Zero-out array elements with matching indices

I'd like to "zero-out" all elements of an n-dimensional array that are in positions with two or more matching indices. In two dimensions, this is effectively np.fill_diagonal() but that function becomes inadequate when a third dimension is considered.

Below is the brute-force version of what I'd like to do. Is there any way to clean this up and make it work in n dimensions?

x = np.ones([3,3,3])

x[:,0,0] = 0
x[0,:,0] = 0
x[0,0,:] = 0

x[:,1,1] = 0
x[1,:,1] = 0
x[1,1,:] = 0

x[:,2,2] = 0
x[2,:,2] = 0
x[2,2,:] = 0

print(x)

Upvotes: 2

Views: 64

Answers (1)

Paul Panzer
Paul Panzer

Reputation: 53029

One way is np.einsum:

>>> a = np.ones((4,4,4), int)
>>> for n in range(3):
...     np.einsum(f"{'iijii'[n:n+3]}->ij", a)[...] = 0
... 
>>> a
array([[[0, 0, 0, 0],
        [0, 0, 1, 1],
        [0, 1, 0, 1],
        [0, 1, 1, 0]],

       [[0, 0, 1, 1],
        [0, 0, 0, 0],
        [1, 0, 0, 1],
        [1, 0, 1, 0]],

       [[0, 1, 0, 1],
        [1, 0, 0, 1],
        [0, 0, 0, 0],
        [1, 1, 0, 0]],

       [[0, 1, 1, 0],
        [1, 0, 1, 0],
        [1, 1, 0, 0],
        [0, 0, 0, 0]]])

General (ND) case:

>>> from string import ascii_lowercase
>>> from itertools import combinations
>>> 
>>> a = np.ones((4,4,4,4), int)
>>> n = a.ndim
>>> ltrs = ascii_lowercase[:n-2]
>>> for I in combinations(range(n), 2):
...     li = iter(ltrs)
...     np.einsum(''.join('z' if k in I else next(li) for k in range(n)) + '->z' + ltrs, a)[...] = 0
... 
>>> a
array([[[[0, 0, 0, 0],
         [0, 0, 0, 0],
         [0, 0, 0, 0],
         [0, 0, 0, 0]],

        [[0, 0, 0, 0],
         [0, 0, 0, 0],
         [0, 0, 0, 1],
         [0, 0, 1, 0]],

        [[0, 0, 0, 0],
         [0, 0, 0, 1],
         [0, 0, 0, 0],
         [0, 1, 0, 0]],

<snip>

Upvotes: 2

Related Questions