NumPy/Pandas: convert array of "steps" into bool mask
Question: I have an array like this:

    
    
    arr = np.array([4, 6, 3, 9, 2, 100, 3, 1, 1, 1, 1])
    

I want to convert it to a bool array like this:

    
    
    [ T, F, F, F, T,   F, T, F, F, T, T]
    # 4, 6, 3, 9, 2, 100, 3, 1, 1, 1, 1
    

I can do it with a loop like this:

    
    
    mask = np.zeros(len(arr), dtype=bool)
    ii = 0
    while ii < len(arr):
        mask[ii] = True
        ii += arr[ii]
    

It's sort of an indirect indexing scheme, where each element in the input
tells us how many subsequent elements are invalid.

How can I do it without using a Python loop, so that it will be fast if the
input array is large? I'm happy to use Pandas too.


Answer: There may be some vectorization trick I'm not thinking of, but if you can use
`numba`, it's well suited for problems like this - this loop should now be
very fast.

    
    
    import numba
    
    @numba.jit(nopython=True)
    def jump_mask(arr):
        mask = np.zeros(len(arr), dtype=np.bool_)
        ii = 0
        while ii < len(arr):
            mask[ii] = True
            ii += arr[ii]
        return mask
    

