CODEWITHSUNDEEP
pythonnumpyscipy
artemka
artemka

Reputation: 59

How do i access all elements of numpy array which indexes satisfy some condition?

For example, I would like to set to zero all elements of a matrix over its counterdiagonal(i + j < n - 1).

I thought about generating a mask, but it would lead to the same problem of accessing such elements in the mask matrix.

What's the best solution?

Upvotes: 2

Views: 93

Answers (1)

Jaime
Jaime

Reputation: 67457

Since your matrix seems to be square, you can use a boolean mask and do:

n = mat.shape[0]
idx = np.arange(n)
mask = idx[:, None] + idx < n - 1
mat[mask] = 0

To understand what's going on:

>>> mat = np.arange(16).reshape(4, 4)
>>> n = 4
>>> idx = np.arange(n)
>>> idx[:, None] + idx
array([[0, 1, 2, 3],
       [1, 2, 3, 4],
       [2, 3, 4, 5],
       [3, 4, 5, 6]])
>>> idx[:, None] + idx < n - 1
array([[ True,  True,  True, False],
       [ True,  True, False, False],
       [ True, False, False, False],
       [False, False, False, False]], dtype=bool)
>>> mat[idx[:, None] + idx < n  -1] = 0
>>> mat
array([[ 0,  0,  0,  3],
       [ 0,  0,  6,  7],
       [ 0,  9, 10, 11],
       [12, 13, 14, 15]])

Upvotes: 4

Related Questions