这是一种矢量化方法 -
def pos_neg_counts(a):
mask = a>0
idx = np.flatnonzero(mask[1:] != mask[:-1])
count = np.concatenate(( [idx[0]+1], idx[1:] - idx[:-1], [a.size-1-idx[-1]] ))
if a[0]<0:
return count[1::2], count[::2] # pos, neg counts
else:
return count[::2], count[1::2] # pos, neg counts
样本运行 -
In [155]: a
Out[155]: array([-3, -2, -1, 1, 2, 3, 4, 5, 6, -5, -4])
In [156]: pos_neg_counts(a)
Out[156]: (array([6]), array([3, 2]))
In [157]: a[0] = 3
In [158]: a
Out[158]: array([ 3, -2, -1, 1, 2, 3, 4, 5, 6, -5, -4])
In [159]: pos_neg_counts(a)
Out[159]: (array([1, 6]), array([2, 2]))
In [160]: a[-1] = 7
In [161]: a
Out[161]: array([ 3, -2, -1, 1, 2, 3, 4, 5, 6, -5, 7])
In [162]: pos_neg_counts(a)
Out[162]: (array([1, 6, 1]), array([2, 1]))
运行时测试
其他方法 -
# @Franz's soln
def split_app(my_array):
negative_index = my_array<0
splits = np.split(negative_index, np.where(np.diff(negative_index))[0]+1)
len_list = [len(i) for i in splits]
return len_list
更大数据集上的时间 -
In [20]: # Setup input array
...: reps = np.random.randint(3,10,(100000))
...: signs = np.ones(len(reps),dtype=int)
...: signs[::2] = -1
...: a = np.repeat(signs, reps)*np.random.randint(1,9,reps.sum())
...:
In [21]: %timeit split_app(a)
10 loops, best of 3: 90.4 ms per loop
In [22]: %timeit pos_neg_counts(a)
100 loops, best of 3: 2.21 ms per loop