在NumPy数组中计算连续的1

Question

[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的向量化解决方案？

- user308827

Answer 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的示例输入数组并重复/平铺它，希望这可以是一个更少机会主义的基准测试-

小案例：

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

- Divakar

---

@rayryeng 奇怪的是，在我的笔记本电脑上（100次重复）Divakar:0.47050797403790057;0.26334572909399867;0.3771821779664606 Paul Panzer:0.6545195570215583;0.4508163968566805;0.17192376195453107 - Paul Panzer

@PaulPanzer 添加了一些不那么机会主义的基准测试。 - Divakar

@Divakar，这很有趣，在我的笔记本电脑上，“loopy one”速度是另一个的两倍，在你的设备上则相反，而对于rayreng，它们似乎几乎一样快。（我的意思是使用a3）。你有什么想法吗？你是直接从shell测试吗？我制作了一个小文件。字节编译会有影响吗？ - Paul Panzer

@PaulPanzer 哈哈，是啊！跟我说说它的情况！ - Divakar

@ConanG 我认为需要进行相当大的修改才能适应二维数组。要不您发一个新问题？ - Divakar

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

- Paul Panzer