Find the index of first non-zero element to the right of given elements in python

Question

I have a 2D numpy.ndarray. Given a list of positions, I want to find the positions of first non-zero elements to the right of the given elements in the same row. Is it possible to vectorize this? I have a huge array and looping is taking too much time.

Eg:

matrix = numpy.array([
    [1, 0, 0, 1, 1], 
    [1, 1, 0, 0, 1], 
    [1, 0, 0, 0, 1], 
    [1, 1, 1, 1, 1], 
    [1, 0, 0, 0, 1]
])
query = numpy.array([[0,2], [2,1], [1,3], [0,1]])

Expected Result:

>> [[0,3], [2,4], [1,4], [0,3]]

Currently I'm doing this using for loops as follows

for query_point in query:
    y, x = query_point
    result_point = numpy.min(numpy.argwhere(self.matrix[y, x + 1:] == 1)) + x + 1
    print(f'{y}, {result_point}')

PS: I also want to find the first non-zero element to the left as well. I guess, the solution to find the right point can be easily tqeaked to find the left point.

Answer 1

If your query array is sufficiently dense, you can reverse the computation: find an array of the same size as matrix that gives the index of the next nonzero element in the same row for each location. Then your problem becomes one of just one of applying query to this index array, which numpy supports directly.

It is actually much easier to find the left index, so let's start with that. We can transform matrix into an array of indices like this:

r, c = np.nonzero(matrix)
left_ind = np.zeros(matrix.shape, dtype=int)
left_ind[r, c] = c

Now you can find the indices of the preceding nonzero element by using np.maximum similarly to how it is done in this answer: https://stackoverflow.com/a/48252024/2988730:

np.maximum.accumulate(left_ind, axis=1, out=left_ind)

Now you can index directly into ind to get the previous nonzero column index:

left_ind[query[:, 0], query[:, 1]]

or

left_ind[tuple(query.T)]

Now to do the same thing with the right index, you need to reverse the array. But then your indices are no longer ascending, and you risk overwriting any zeros you have in the first column. To solve that, in addition to just reversing the array, you need to reverse the order of the indices:

right_ind = np.zeros(matrix.shape, dtype=int)
right_ind[r, c] = matrix.shape[1] - c

You can use any number larger than matrix.shape[1] as your constant as well. The important thing is that the reversed indices all come out greater than zero so np.maximum.accumulate overwrites the zeros. Now you can use np.maximum.accumulate in the same way on the reversed array:

right_ind = matrix.shape[1] - np.maximum.accumulate(right_ind[:, ::-1], axis=1)[:, ::-1]

In this case, I would recommend against using out=right_ind, since right_ind[:, ::-1] is a view into the same buffer. The operation is buffered, but if your line size is big enough, you may overwrite data unintentionally.

Now you can index the array in the same way as before:

right_ind[(*query.T,)]

In both cases, you need to stack with the first column of query, since that's the row key:

>>> row, col = query.T
>>> np.stack((row, left_ind[row, col]), -1)
array([[0, 0],
       [2, 0],
       [1, 1],
       [0, 0]])
>>> np.stack((row, right_ind[row, col]), -1)
array([[0, 3],
       [2, 4],
       [1, 4],
       [0, 3]])
>>> np.stack((row, left_ind[row, col], right_ind[row, col]), -1)
array([[0, 0, 3],
       [2, 0, 4],
       [1, 1, 4],
       [0, 0, 3]])

If you plan on sampling most of the rows in the array, either at once, or throughout your program, this will help you speed things up. If, on the other hand, you only need to access a small subset, you can apply this technique only to the rows you need.