Find consecutive sequences based on Boolean array

Question

I'm trying to extract sequences from an array b for which a boolean array a is used as index (len(a) >= len(b), but (a==True).sum() == len(b), i.e. there are only as many values true in a than there are elements in b). The sequences should be represented in the result as start and end index of a where a[i] is true and for which there are consecutive values.

For instance, for the following arrays of a and b

a = np.asarray([True, True, False, False, False, True, True, True, False])
b = [1, 2, 3, 4, 5]

the result should be [((0, 1), [1, 2]), ((5, 7), [3, 4, 5])], so as many elements in the array as there are true-sequences. Each true sequence should contain the start and end index from a and the values these relate to from b).

So for the above:

[
 ((0, 1), [1, 2]),   # first true sequence: starting at index=0 (in a), ending at index=1, mapping to the values [1, 2] in b

 ((5, 7), [3, 4, 5]) # second true sequence: starting at index=5, ending at index=7, with values in b=[3, 4, 5]
]

How can this be done efficiently in numpy?

Answer 1

Here's one NumPy based one inspired by this post -

def func1(a,b):
    # "Enclose" mask with sentients to catch shifts later on
    mask = np.r_[False,a,False]

    # Get the shifting indices
    idx = np.flatnonzero(mask[1:] != mask[:-1])

    s0,s1 = idx[::2], idx[1::2]
    idx_b = np.r_[0,(s1-s0).cumsum()]
    out = []
    for (i,j,k,l) in zip(s0,s1-1,idx_b[:-1],idx_b[1:]):
        out.append(((i, j), b[k:l]))
    return out

Sample run -

In [104]: a
Out[104]: array([ True,  True, False, False, False,  True,  True,  True, False])

In [105]: b
Out[105]: [1, 2, 3, 4, 5]

In [106]: func1(a,b)
Out[106]: [((0, 1), [1, 2]), ((5, 7), [3, 4, 5])]

Timings -

In [156]: # Using given sample data and tiling it 1000x
     ...: a = np.asarray([True, True, False, False, False, True, True, True, False])
     ...: b = [1, 2, 3, 4, 5]
     ...: a = np.tile(a,1000)
     ...: b = np.tile(b,1000)

# @Chris's soln
In [157]: %%timeit
     ...: res = []
     ...: gen = (i for i in b)
     ...: for k, g in itertools.groupby(enumerate(a), lambda x:x[1]):
     ...:     if k:
     ...:         ind, bools = list(zip(*g))
     ...:         res.append((ind[0::len(ind)-1], list(itertools.islice(gen, len(bools)))))
100 loops, best of 3: 13.8 ms per loop

In [158]: %timeit func1(a,b)
1000 loops, best of 3: 1.29 ms per loop