python - Pandas: assign an index to each group identified by groupby

Question

When using groupby(), how can I create a DataFrame with a new column containing an index of the group number, similar to dplyr::group_indices in R. For example, if I have

>>> df=pd.DataFrame({'a':[1,1,1,2,2,2],'b':[1,1,2,1,1,2]})
>>> df
   a  b
0  1  1
1  1  1
2  1  2
3  2  1
4  2  1
5  2  2

How can I get a DataFrame like

   a  b  idx
0  1  1  1
1  1  1  1
2  1  2  2
3  2  1  3
4  2  1  3
5  2  2  4

(the order of the idx indexes doesn't matter)

Answer 1

Here is the solution using ngroup (available as of pandas 0.20.2) from a comment above by Constantino, for those still looking for this function (the equivalent of dplyr::group_indices in R, or egen group() in Stata if you were trying to search with those keywords like me). This is also about 25% faster than the solution given by maxliving according to my own timing.

>>> import pandas as pd
>>> df = pd.DataFrame({'a':[1,1,1,2,2,2],'b':[1,1,2,1,1,2]})
>>> df['idx'] = df.groupby(['a', 'b']).ngroup()
>>> df
   a  b  idx
0  1  1    0
1  1  1    0
2  1  2    1
3  2  1    2
4  2  1    2
5  2  2    3

>>> %timeit df['idx'] = create_index_usingduplicated(df, grouping_cols=['a', 'b'])
1.83 ms ± 67.2 μs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
>>> %timeit df['idx'] = df.groupby(['a', 'b']).ngroup()
1.38 ms ± 30 μs per loop (mean ± std. dev. of 7 runs, 1000 loops each)