Home > Net >  Fastest way to slice a numpy 1d array based on consecutive True values
Fastest way to slice a numpy 1d array based on consecutive True values

Time:09-02

I have a numpy 1d array like this:

a = np.array([False, True, True, True, False, True, True, False])

I want to find the index of the first True value in the array (1 in this case) and the last index of the first consecutive True sequence (3 in this case).

I actually can imagine a lot of solutions for this problem, but I'm looking for the cleanest way to solve this task. For simplicity assume that there is almost one True in the array

some other examples could be useful (expected output as #comment):

[False, True, True] # 1, 2
[True, True, True, False, True, True, False] # 0, 2
[False, True, True, True] # 1, 3
[False, True, False, True] # 1, 1

CodePudding user response:

My method would be:

  1. use np.where to find indexes of all True values
  2. add 1 and use np.setdiff1d to find the first False after the first True
  3. return the first idx found in step (1), and the predecessor of the idx found in step (2).
def slice_true_values(a):
    w = np.where(a)[0]
    d = np.setdiff1d(w   1, w)
    return w[0], d[0] - 1

Trying to explain the logic: when we add 1 to each index of the where array, all consecutive indexes will be present in both the original and the 1 array, except for the last one.

  • Related