计算具有多个值的数组中的连续值 numpy/pandas

Question

I checked this question and others on SO but the trick is always summing True or False values.

My case is the following array:

arr = [1,2,3,3,4,5,6,1,1,1,5,5,8,8,8,9,4,4,4]

I want to get for each member of the array the length of the "current" streak of repeated value.

For the example above I would like to get:

res = [1,1,1,2,1,1,1,1,2,3,1,2,1,2,3,1,1,2,3]

I could write a dumb loop but is there a clever or already built-in way to do this in numpy/pandas?

Answer 1

A pandas way for input array arr would be -

In [55]: I = np.r_[False,arr[:-1]!=arr[1:]].cumsum()

In [56]: df = pd.DataFrame({'ids':I,'val':np.ones(len(arr),dtype=int)})

In [57]: df.groupby('ids')[['val']].cumsum().values.ravel()
Out[57]: array([1, 1, 1, 2, 1, 1, 1, 1, 2, 3, 1, 2, 1, 2, 3, 1, 1, 2, 3])

Another with a custom NumPy implementation to create ranges based on interval lengths/sizes - intervaled_ranges -

In [91]: m = np.r_[True,arr[:-1]!=arr[1:],True]

In [92]: intervaled_ranges(np.diff(np.flatnonzero(m)),start=1)
Out[92]: array([1, 1, 1, 2, 1, 1, 1, 1, 2, 3, 1, 2, 1, 2, 3, 1, 1, 2, 3])

Answer 2

You can also use pd.Series and groupby :

s = pd.Series([1,2,3,3,4,5,6,1,1,1,5,5,8,8,8,9,4,4,4])

print (s.groupby((s!=s.shift()).cumsum()).cumcount() + 1)
#
[1, 1, 1, 2, 1, 1, 1, 1, 2, 3, 1, 2, 1, 2, 3, 1, 1, 2, 3]