我有一个 1 和 0 的 Numpy 一维数组。例如
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”解决方案吗？
【问题讨论】：
这是一种矢量化方法 -
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
【讨论】：
使用来自itertools 的groupby
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)]
【讨论】：