10

I have a numpy array of zeros and ones:

y=[1,1,1,0,0,0,0,0,1,1,0,0,0,0,0,0,1,1,1,1]

I want to calculate the indices of groups of ones (or zeros). So for the above example the result for groups of ones should be something similar to:

result=[(0,2), (8,9), (16,19)]

(How) Can I do that with numpy? I found nothing like a group-by function.

I experimented around with np.ediff1d, but couldn't figure out a good solution. Not that the array may or may not begin/end with a group of ones:

import numpy as np

y = [1,1,1,0,0,0,0,0,1,1,0,0,0,0,0,0,1,1,1,1]
mask = np.ediff1d(y)
starts = np.where(mask > 0)
ends = np.where(mask < 0)

I also found a partial solution here: Find index where elements change value numpy

But that one only gives me the indices where the values change.

tiefenauer
  • 669
  • 7
  • 19

1 Answers1

11

We can do something like this that works for any generic array -

def islandinfo(y, trigger_val, stopind_inclusive=True):
    # Setup "sentients" on either sides to make sure we have setup
    # "ramps" to catch the start and stop for the edge islands
    # (left-most and right-most islands) respectively
    y_ext = np.r_[False,y==trigger_val, False]

    # Get indices of shifts, which represent the start and stop indices
    idx = np.flatnonzero(y_ext[:-1] != y_ext[1:])

    # Lengths of islands if needed
    lens = idx[1::2] - idx[:-1:2]

    # Using a stepsize of 2 would get us start and stop indices for each island
    return list(zip(idx[:-1:2], idx[1::2]-int(stopind_inclusive))), lens

Sample run -

In [320]: y
Out[320]: array([1, 1, 1, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1])

In [321]: islandinfo(y, trigger_val=1)[0]
Out[321]: [(0, 2), (8, 9), (16, 19)]

In [322]: islandinfo(y, trigger_val=0)[0]
Out[322]: [(3, 7), (10, 15)]

Alternatively, we can use diff to get the sliced comparisons and then simply reshape with 2 columns to replace the step-sized slicing to give ourselves a one-liner -

In [300]: np.flatnonzero(np.diff(np.r_[0,y,0])!=0).reshape(-1,2) - [0,1]
Out[300]: 
array([[ 0,  2],
       [ 8,  9],
       [16, 19]])
Divakar
  • 218,885
  • 19
  • 262
  • 358
  • can you explain more on how the faster method works? – R.yan May 03 '18 at 09:36
  • I used the second solution and it works. I'm marking this post as the answer although I'm probably going to have a hard time understanding this again in a few weeks :) – tiefenauer May 03 '18 at 11:06
  • 1
    With Python 3, in the last line of `islandof`, `zip(...)` should be replaced with `list(zip(...))`. This is because `zip` returns an iterator in Python 3. – JCOidl Apr 08 '20 at 21:23
  • I used `%timeit` and at least on my machine the `np.concatenate(([0], y, [0]))` is much quicker than `np.r_[0,y,0]`. – acmpo6ou Mar 05 '23 at 20:07