Creating index array of next valid boolean value in numpy

user_14484091 0
Viewed 103

I have a one dimensional numpy boolean-array and my goal is to get an integer-array of the same dimensionality in which each entry states the next valid index of the boolean array. Here an example:

import numpy as np
b=np.asarray((True, False, False, True, False),dtype=bool)
#It should be assigned such that in the end
#i=array([ 0,  3,  3,  3, -1])

Is there any combination numpy functions which achieves this in linear runtime and without the use of loops?

numpy arrays python
2 Answers

This is one way I can think of doing that:

import numpy as np
b = np.asarray((True, False, False, True, False), dtype=bool)
idx = np.r_[np.where(b)[0], -1]
res = idx[np.cumsum(np.r_[False, b[:-1]])]
print(res)
# [ 0  3  3  3 -1]

If b is very big, you can save one concatenation with:

import numpy as np
b = np.asarray((True, False, False, True, False), dtype=bool)
idx = np.r_[np.where(b)[0], -1]
c = np.zeros(len(b), dtype=np.int32)
np.cumsum(b[:-1], out=c[1:])
res = idx[c]
print(res)
# [ 0  3  3  3 -1]
user_17827920

Here's my idea to solve that problem:

import numpy as np


def next_valid_index(b):

    b = np.asarray(b, dtype=bool)

    # Initialize output
    i = np.int32(b)

    # Set True elements
    ind = np.argwhere(b)
    i[ind] = ind

    # Set False elements
    i[np.where(1 - b)[0]] = \
        ind[np.argmax(np.argwhere(1 - b) < ind.T, axis=1)].squeeze()
    i[np.where((b == False) & (i == 0))[0]] = -1

    return i


B = (True, False, False, True, False)
I = next_valid_index(B)
print(B, '\n', I)

B = (True, False, False, True, False, True, True, False)
I = next_valid_index(B)
print(B, '\n', I)

B = (False, False, False, True, False, True, True, False, True)
I = next_valid_index(B)
print(B, '\n', I)

Examples and corresponding outputs:

(True, False, False, True, False) 
 [ 0  3  3  3 -1]
(True, False, False, True, False, True, True, False) 
 [ 0  3  3  3  5  5  6 -1]
(False, False, False, True, False, True, True, False, True) 
 [3 3 3 3 5 5 6 8 8]

----------------------------------------
System information
----------------------------------------
Platform:    Windows-10-10.0.16299-SP0
Python:      3.8.5
NumPy:       1.19.2
----------------------------------------
user_110899320
Related