CODEWITHSUNDEEP
pythonarraysnumpyindexing
user8991539
user8991539

Reputation:

NumPy where as range indices

I have a numpy array like that:

a = numpy.array([1, 1, 1, 0, 0, 0, 1, 1, 1, 0, 0, 0])

The goal is to find the ranges of zeros and ones, start and end indices. I want to use this ranges in another numpy array which contains time stamps to find out how much time each zero phase takes. Something like that:

dur = numpy.diff(time[start idx, end idx])

However, numpy where gives me all indices:

numpy.where(a==0)
(array([ 3,  4,  5,  9, 10, 11], dtype=int64),)

I would need only start and end idx of each zero phase, like [[3,5],[9,11]]. How can I achieve that?

Upvotes: 1

Views: 511

Answers (1)

Divakar
Divakar

Reputation: 221514

Here's one approach -

def start_stop(a, val=0):
    n = np.concatenate(([False], a==val,[False]))
    idx = np.flatnonzero(np.diff(n))
    # or idx = np.flatnonzero(n[1:] != n[:-1])
    return idx[::2], idx[1::2]-1

Shorter way -

def start_stop_v2(a, val=0):
    idx = np.flatnonzero(np.diff(np.r_[0,a==val,0]))
    return idx[::2], idx[1::2]-1

One-liner -

np.flatnonzero(np.diff(np.r_[0,a==0,0])).reshape(-1,2) - [0,1]

Sample run -

In [324]: a
Out[324]: array([1, 1, 1, 0, 0, 0, 1, 1, 1, 0, 0, 0])

In [325]: start_stop(a, val=0)
Out[325]: (array([3, 9]), array([ 5, 11]))

In [326]: start_stop_v2(a, val=0)
Out[326]: (array([3, 9]), array([ 5, 11]))

In [327]: np.flatnonzero(np.diff(np.r_[0,a==0,0])).reshape(-1,2) - [0,1]
Out[327]: 
array([[ 3,  5],
       [ 9, 11]])

Re-using of np.where(a==val)

To solve it re-using result of np.where(a==val) -

In [388]: idx = numpy.where(a==0)[0]

In [389]: mask = np.r_[True,np.diff(idx)!=1,True]

In [390]: idx[mask[:-1]] # starts
Out[390]: array([3, 9])

In [391]: idx[mask[1:]] # stops
Out[391]: array([ 5, 11])

Upvotes: 1

Related Questions