python - 根据 bool 数组查找连续序列

Question

我正在尝试从数组 b 中提取序列，其中 bool 数组 a 用作索引 (len(a) >= len( b)，但是 (a==True).sum() == len(b)，即 a 中的 true 值只比b 中有元素)。序列应在结果中表示为 a 的开始和结束索引，其中 a[i] 为真且有连续值。

例如，对于下面的a和b

数组

a = np.asarray([True, True, False, False, False, True, True, True, False])
b = [1, 2, 3, 4, 5]

结果应该是[((0, 1), [1, 2]), ((5, 7), [3, 4, 5])]，所以尽可能多的元素在数组中，因为有真实序列。每个真实序列都应包含 a 的开始和结束索引以及这些与 b 相关的值。

所以对于上面的内容:

[
 ((0, 1), [1, 2]),   # first true sequence: starting at index=0 (in a), ending at index=1, mapping to the values [1, 2] in b

 ((5, 7), [3, 4, 5]) # second true sequence: starting at index=5, ending at index=7, with values in b=[3, 4, 5]
]

如何在 numpy 中高效地完成这项工作？

Answer 1

这是一个基于 NumPy 的灵感来自 this post 的一个-

def func1(a,b):
    # "Enclose" mask with sentients to catch shifts later on
    mask = np.r_[False,a,False]

    # Get the shifting indices
    idx = np.flatnonzero(mask[1:] != mask[:-1])

    s0,s1 = idx[::2], idx[1::2]
    idx_b = np.r_[0,(s1-s0).cumsum()]
    out = []
    for (i,j,k,l) in zip(s0,s1-1,idx_b[:-1],idx_b[1:]):
        out.append(((i, j), b[k:l]))
    return out

sample 运行-

In [104]: a
Out[104]: array([ True,  True, False, False, False,  True,  True,  True, False])

In [105]: b
Out[105]: [1, 2, 3, 4, 5]

In [106]: func1(a,b)
Out[106]: [((0, 1), [1, 2]), ((5, 7), [3, 4, 5])]

时间 -

In [156]: # Using given sample data and tiling it 1000x
     ...: a = np.asarray([True, True, False, False, False, True, True, True, False])
     ...: b = [1, 2, 3, 4, 5]
     ...: a = np.tile(a,1000)
     ...: b = np.tile(b,1000)

# @Chris's soln
In [157]: %%timeit
     ...: res = []
     ...: gen = (i for i in b)
     ...: for k, g in itertools.groupby(enumerate(a), lambda x:x[1]):
     ...:     if k:
     ...:         ind, bools = list(zip(*g))
     ...:         res.append((ind[0::len(ind)-1], list(itertools.islice(gen, len(bools)))))
100 loops, best of 3: 13.8 ms per loop

In [158]: %timeit func1(a,b)
1000 loops, best of 3: 1.29 ms per loop