。例如
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])
I want to count the continuous 0s and 1s in the array and output something like this
[1,3,7,1,1,2,3,2,2]
我做的是atm
np.diff(np.where(np.abs(np.diff(a)) == 1)[0])
它输出
array([3, 7, 1, 1, 2, 3, 2])
正如你所看到的,它缺少第一个计数1。
我试过
np.split
然后得到每个段的大小,但似乎并不乐观。有更优雅的“蟒蛇”解决方案吗?
最佳答案
-
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])
连接速度更快-
np.diff(np.flatnonzero(np.concatenate(([True], a[1:]!= a[:-1], [True] ))))
运行时测试
对于设置,让我们创建一个包含
boolean
和0s
孤岛的更大数据集,对于给定样本的公平基准,让孤岛长度在1s
和1
之间变化-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