计算 NumPy 数组中的连续 1答案

【问题标题】：Counting consecutive 1's in NumPy array计算 NumPy 数组中的连续 1
【发布时间】：2017-02-09 05:37:13
【问题描述】：

[1, 1, 1, 0, 0, 0, 1, 1, 0, 0]

我有一个像上面一样由 0 和 1 组成的 NumPy 数组。如何添加所有连续的 1，如下所示？每当我遇到 0 时，我都会重置。

[1, 2, 3, 0, 0, 0, 1, 2, 0, 0]

我可以使用 for 循环来做到这一点，但有没有使用 NumPy 的矢量化解决方案？

【问题讨论】：

标签： python numpy cumsum

【解决方案1】：

这是一种矢量化方法 -

def island_cumsum_vectorized(a):
    a_ext = np.concatenate(( [0], a, [0] ))
    idx = np.flatnonzero(a_ext[1:] != a_ext[:-1])
    a_ext[1:][idx[1::2]] = idx[::2] - idx[1::2]
    return a_ext.cumsum()[1:-1]

示例运行 -

In [91]: a = np.array([1, 1, 1, 0, 0, 0, 1, 1, 0, 0])

In [92]: island_cumsum_vectorized(a)
Out[92]: array([1, 2, 3, 0, 0, 0, 1, 2, 0, 0])

In [93]: a = np.array([0, 1, 1, 1, 1, 0, 0, 0, 1, 1, 0, 0, 1])

In [94]: island_cumsum_vectorized(a)
Out[94]: array([0, 1, 2, 3, 4, 0, 0, 0, 1, 2, 0, 0, 1])

运行时测试

对于时间安排，我会使用 OP 的示例输入数组并重复/平铺它，希望这应该是 less opportunistic benchmark -

小号：

In [16]: a = np.array([1, 1, 1, 0, 0, 0, 1, 1, 0, 0])

In [17]: a = np.tile(a,10)  # Repeat OP's data 10 times

# @Paul Panzer's solution
In [18]: %timeit np.concatenate([np.cumsum(c) if c[0] == 1 else c for c in np.split(a, 1 + np.where(np.diff(a))[0])])
10000 loops, best of 3: 73.4 µs per loop

In [19]: %timeit island_cumsum_vectorized(a)
100000 loops, best of 3: 8.65 µs per loop

更大的情况：

In [20]: a = np.array([1, 1, 1, 0, 0, 0, 1, 1, 0, 0])

In [21]: a = np.tile(a,1000)  # Repeat OP's data 1000 times

# @Paul Panzer's solution
In [22]: %timeit np.concatenate([np.cumsum(c) if c[0] == 1 else c for c in np.split(a, 1 + np.where(np.diff(a))[0])])
100 loops, best of 3: 6.52 ms per loop

In [23]: %timeit island_cumsum_vectorized(a)
10000 loops, best of 3: 49.7 µs per loop

不，我想要一个非常大的箱子：

In [24]: a = np.array([1, 1, 1, 0, 0, 0, 1, 1, 0, 0])

In [25]: a = np.tile(a,100000)  # Repeat OP's data 100000 times

# @Paul Panzer's solution
In [26]: %timeit np.concatenate([np.cumsum(c) if c[0] == 1 else c for c in np.split(a, 1 + np.where(np.diff(a))[0])])
1 loops, best of 3: 725 ms per loop

In [27]: %timeit island_cumsum_vectorized(a)
100 loops, best of 3: 7.28 ms per loop

【讨论】：

对这些进行基准测试：a1 = np.repeat(np.arange(1000)%2, np.random.randint(1, 1000, (1000,))); a2 = np.repeat(np.arange(100)%2, np.random.randint(1, 10000, (100,))); a3 = np.repeat(np.random.random((100,)) < 0.2, np.random.randint(1, 10000, (100,))) ;-)
@PaulPanzer 我决定自己根据您的建议输入运行所有这些基准测试。 Divakar 的实现仍然更快：P... 大约 2 倍。
@rayreng 也在最后一个？但我承认他的更好；我只是反对机会主义的基准;-)
@PaulPanzer :D。我已经赞成你的，因为它非常好，而且在一条线上。我刚刚跑的最后一个......它只是略微领先。
@rayryeng 奇怪，在我的笔记本电脑上（100 次重复）Divakar:0.47050797403790057;0.26334572909399867;0.3771821779664606Paul Panzer:0.6545195570215583;0.4508163968566805;0.17192376195453107

【解决方案2】：

如果列表理解是可以接受的

np.concatenate([np.cumsum(c) if c[0] == 1 else c for c in np.split(a, 1 + np.where(np.diff(a))[0])])

【讨论】：