NumPy作为范围索引

时间:2017-11-28 12:25:41

标签: python arrays numpy indexing

我有一个像那样的numpy数组:

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

目标是找到零和1的范围,开始和结束索引。我想在另一个包含时间戳的numpy数组中使用这个范围,以找出每个零阶段需要多少时间。这样的事情:

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

然而,numpy在哪里给我所有索引:

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

我只需要每个零阶段的开始和结束idx,如[[3,5],[9,11]]。我怎样才能做到这一点?

1 个答案:

答案 0 :(得分:1)

这是一种方法 -

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

缩短方式 -

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]

示例运行 -

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]])

重复使用np.where(a==val)

要解决它重复使用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])