Masking few non-zero elements of certain rows of a matrix
I have a 3*3 matrix with 1s and 0s, A = [[1,0,1],[0,1,1],[1,0,0]] and an array indicating the limit on the row sum, B = [1,2,1]. I want to find rows of A whose sum exceeds the corresponding value in B, and set the non-zero elements of A to zero to ensure that the sum matches with B. Finding the rows of A that exceed the sum is easy, however masking the elements to adjust the sum is what I need help with.How can this be achieved (want to scale it to larger matrices and tensors)?

I would do smth like this:
import numpy as np

A = np.array([[1,0,1],[0,1,1],[1,0,0]])
B = np.array([1,2,1])

# a cumulative sum of each row will tell you how many 
# ones were in that row up to each point.
A_cs = np.cumsum(A, axis = 1)
# theresholding according to the sum vector will let 
# you know where you should start omitting values since 
# at that point the sum of the row exceeds its limit.
A_th = A_cs &gt; B[:, None]
# then you can use the boolean array to create a new 
# array where the appropriate values in the original 
# array are set to zero to reduce the row sum.
A_nw = A * (1 - A_th)

output:
A_nw = 
[[1 0 0]
 [0 1 1]
 [1 0 0]]

Unrelated note:
The following note is here to help OP better their dev-related search skills.
I can answer some questions instantaneously, but this was not one of them. I'm telling you this because I reached the answer through a simple google search for &quot;python find the i th non zero element in each row&quot;, which led me to this post which in turn led me very quickly to an answer. You don't have to try to be a better, more independet code writer. But if you want to, know that you can.
