Question

我有一个Numpy一维数组1和0.例如

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

我想计算数组中的连续0和1并输出类似这样的内容

[1,3,7,1,1,2,3,2,2]

我做的是什么

np.diff(np.where(np.abs(np.diff(a)) == 1)[0])

并输出

array([3, 7, 1, 1, 2, 3, 2])

你可以看到它缺少第一个计数1。

我已尝试np.split，然后获取每个细分的大小，但似乎并不乐观。

是否有更优雅的“pythonic”解决方案？

Answer 1

这是一种矢量化方法 -

np.diff(np.r_[0,np.flatnonzero(np.diff(a))+1,a.size])

示例运行 -

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

In [209]: np.diff(np.r_[0,np.flatnonzero(np.diff(a))+1,a.size])
Out[209]: array([1, 3, 7, 1, 1, 2, 3, 2, 2])

boolean连接速度更快 -

np.diff(np.flatnonzero(np.concatenate(([True], a[1:]!= a[:-1], [True] ))))

运行时测试

对于设置，让我们创建一个包含0s和1s孤岛的更大数据集，并且对于与给定样本一样公平的基准测试，让岛长度在1和7 -

In [257]: n = 100000 # thus would create 100000 pair of islands

In [258]: a = np.repeat(np.arange(n)%2, np.random.randint(1,7,(n)))

# Approach #1 proposed in this post
In [259]: %timeit np.diff(np.r_[0,np.flatnonzero(np.diff(a))+1,a.size])
100 loops, best of 3: 2.13 ms per loop

# Approach #2 proposed in this post
In [260]: %timeit np.diff(np.flatnonzero(np.concatenate(([True], a[1:]!= a[:-1], [True] ))))
1000 loops, best of 3: 1.21 ms per loop

# @Vineet Jain's soln    
In [261]: %timeit [ sum(1 for i in g) for k,g in groupby(a)]
10 loops, best of 3: 61.3 ms per loop

Answer 2

使用groupby

中的itertools

from itertools import groupby
a = np.array([0,1,1,1,0,0,0,0,0,0,0,1,0,1,1,0,0,0,1,1,0,0])
grouped_a = [ sum(1 for i in g) for k,g in groupby(a)]

如何计算numpy中的连续数字

2 个答案: