CODEWITHSUNDEEP
pythonnumpymatrixindexingvectorization
Leguan3000
Leguan3000

Reputation: 133

Set numpy matrix elements to zero if varying row index is exceeded

I have a quite large m times n numpy matrix M filled with non-zero values and an array x of length m, where each entry indicates the row index, after which the matrix elements should be set to zero. So for example, if n=5 and x[i]=3, then the i-th row of the matrix be set to [M_i1, M_i2, M_i3, 0, 0].

If all entries of x had the same value k, I could simply use slicing with something like M[:,k:]=0, but I could not figure out an efficient way to this with different values for each row without looping over all rows and use slicing for each row.

I thougt about creating a matrix that looks like [[1]*x[1] + [0]*(n-x[1]),...,[1]*x[m] + [0]*(n-x[m])] and use it for boolean indexing but also don't know how to create this without looping.

The non-vectorized solution looks like this:

for i in range(m):
    if x[i] < n:
        M[i,x[i]:] = 0

with example input

M = np.array([[1,2,3],[4,5,6]])
m, n =  2, 3
x = np.array([1,2])

and output

array([[1, 0, 0],
       [4, 5, 0]])

Does anyone have a vectorized solution for this problem?

Thank you very much!

Upvotes: 3

Views: 1104

Answers (2)

Mad Physicist
Mad Physicist

Reputation: 114528

This looks like a mask-smearing exercise. At each row, you want to smear starting with the element at np.minimum(x[row], n):

mask = np.zeros(M.shape, bool)
mask[np.flatnonzero(x < n), x[x < n]] = True
M[np.cumsum(mask, axis=1, dtype=bool)] = 0

Upvotes: 0

Ehsan
Ehsan

Reputation: 12417

You can use multi-dimensional boolean indexing:

M[x[:,None]<=np.arange(M.shape[1])] = 0

example:

M = [[7, 8, 4, 2, 3, 9, 1, 8, 4, 3],
     [2, 1, 6, 1, 5, 2, 2, 2, 9, 2],
     [6, 1, 6, 8, 4, 3, 6, 9, 2, 6],
     [5, 4, 0, 8, 3, 0, 0, 1, 8, 7],
     [8, 7, 8, 8, 9, 2, 0, 8, 0, 2]]
x = [4, 4, 0, 6, 2]

output:

      [[7, 8, 4, 2, 0, 0, 0, 0, 0, 0],
       [2, 1, 6, 1, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [5, 4, 0, 8, 3, 0, 0, 0, 0, 0],
       [8, 7, 0, 0, 0, 0, 0, 0, 0, 0]]

Upvotes: 1

Related Questions