pandas increment count column based on value of another column

Question

I have a df that is the result of a join:

   ID   count
0  A   30
1  A   30
2  B   5
3  C   44
4  C   44
5  C   44

I would like to be able to iterate the count column based on the ID column. Here is an example of the desired result:

   ID   count
0  A   30
1  A   31
2  B   5
3  C   44
4  C   45
5  C   46

I know there are non-pythonic ways to do this via loops, but I am wondering if there is a more intelligent (and time efficient, as this table is large) way to do this.

Answer 1

Transform the group to get a cumulative count and add it to count, eg:

df['count'] += df.groupby('ID')['count'].cumcount()

Gives you:

  ID  count
0  A     30
1  A     31
2  B      5
3  C     44
4  C     45
5  C     46