Creating index array of next valid boolean value in numpy

Question

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?

If this is not possible this will be my fallback solution. The dataset which I have is relatively large so I'm just curious if there is already a nice way to do this in numpy with less runtime. — Paul Stahr
– Paul Stahr, Commented Oct 20, 2020 at 9:25
On a 100000 element random boolean array my old for loop-approach took 19.6ms, jdehesa's 1.83ms and HansHirse's 2.15s per iteration. — Paul Stahr
– Paul Stahr, Commented Oct 20, 2020 at 10:26

javidcf · Accepted Answer · 2020-10-20 09:38:54Z

4

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]

answered Oct 20, 2020 at 9:38

javidcf

59.9k7 gold badges87 silver badges134 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

HansHirse · Accepted Answer · 2020-10-20 09:43:09Z

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
----------------------------------------

Collectives™ on Stack Overflow

Creating index array of next valid boolean value in numpy

2 Answers 2

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related