groupby consecutive identical values in pandas dataframe and cumulative count of the number of occurences

Viewed 138

I have a problem where I would like to count the number of times the current value has not changed in a dataframe over rolling periods.

For example:

df = pd.DataFrame({'col':list('aaaabbab')})

would somehow give output of

0
1
2
3
0
1
0
0

I have been trying something along the following

df['col'] = df['col'] == df['col'].shift(1)
df.rolling(window=3).sum().reset_index(drop=True, level=0)

I have added in the rolling as I will want to look at the full data set in terms of rolling periods but even without having it over rolling periods I can not quite figure out the logic.

I am not sure if I am missing something simple or this may not be possible using shift

1 Answers

You need to generate a grouper for the change in values. For this compare each value with the previous one and apply a cumsum. This gives you groups in the itertools.groupby style ([1, 1, 1, 1, 2, 2, 3, 4]), finally group and apply a cumcount.

df['count'] = (df.groupby(df['col'].ne(df['col'].shift()).cumsum())
                 .cumcount()
              )

output:

  col  count
0   a      0
1   a      1
2   a      2
3   a      3
4   b      0
5   b      1
6   a      0
7   b      0

edit: for fun here is a solution using itertools (much faster):

from itertools import groupby, chain
df['count'] = list(chain(*(list(range(len(list(g)))) 
                           for _,g in  groupby(df['col']))))

NB. this runs much faster (88 µs vs 707 µs on the provided example)

Related