numpy matrix of distances from first non null element in the row

user_3623123 0
Viewed 43

I have a numpy array where in each row there is at most 1 non zero element. I need to efficiently build a matrix where in the element (i,j) there is the horizontal distance from the non zero element in that row.

For istante assume that element (10,6) is 1, in the matrix i need to build element (10,6) is 0, element (10,5) is 1 element (10,7) is -1 etc.

I need this algorithm to be efficient. I’d like to avoid for loops and uso broadcastable functions like np.roll

broadcast numpy python
1 Answers

If there is a one in every row, you can do it with a simple broadcasted subtraction:

np.arange(arr.shape[0]) - arr.argmax(1)[:, None]
user_29887300
Related