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.

Answer 2

I came up with a solution to get both your wanted indices, i.e. to the left and to the right from the indicated position.

First define the following function, to get the row number and both indices:

def inds(r, c, arr):
    ind = np.nonzero(arr[r])[0]
    indSlice = ind[ind < c]
    iLeft = indSlice[-1] if indSlice.size > 0 else None
    indSlice = ind[ind > c]
    iRight = indSlice[0] if indSlice.size > 0 else None
    return r, iLeft, iRight

Parameters:

• r and c are row number (in the source array) and the "starting" index in this row,
• arr is the array to look in (matrix will be passed here).

Then define the vectorized version of this function:

indsVec = np.vectorize(inds, excluded=['arr'])

And to get the result, run:

result = np.vstack(indsVec(query[:, 0], query[:, 1], arr=matrix)).T

The result is:

array([[0, 0, 3],
       [2, 0, 4],
       [1, 1, 4],
       [0, 0, 3]], dtype=int64)

Your expected result is the left and right column (row number and the index of first non-zero element after the "starting" position.

The middle column is the index of last non-zero element before the "starting" position.

This solution is resistant to "non-existing" case (if there are no any "before" or "after" non-zero element). In such case the respective index is returned as None.