How to label groups of pairs in pandas?

Question

I have this dataframe:

>>> df = pd.DataFrame({'A': [1, 2, 1, np.nan, 2, 2, 2], 'B': [2, 1, 2, 2.0, 1, 1, 2]})
>>> df
     A    B
0  1.0  2.0
1  2.0  1.0
2  1.0  2.0
3  NaN  2.0
4  2.0  1.0
5  2.0  1.0
6  2.0  2.0

I need to identify the groups of pairs (A,B) on a third column "group id", to get something like this:

>>> df
     A    B  grup id                        explanation
0  1.0  2.0      1.0  <- group (1.0, 2.0), first group 
1  2.0  1.0      2.0  <- group (2.0, 1.0), second group
2  1.0  2.0      1.0  <- group (1.0, 2.0), first group 
3  NaN  2.0      NaN  <- invalid group                 
4  2.0  1.0      2.0  <- group (2.0, 1.0), second group
5  2.0  1.0      2.0  <- group (2.0, 1.0), second group
6  2.0  2.0      3.0  <- group (2.0, 2.0), third group 

How can I do this efficiently in pandas?

One idea is to first build a combined column (A,B), then identify the unique values in that column and map them back to my dataframe. But I suspect that a groupby() approach would be faster (and more elegant).

I tried this:

>>> df.groupby(['A','B']).count()
Empty DataFrame
Columns: []
Index: [(1.0, 2.0), (2.0, 1.0), (2.0, 2.0)]

So the index of this groupby() lists all the groups I need. But then how to count them and map them back to my dataframe?

Answer 1

You can use GroupBy.ngroup (pandas 0.20.2+):

print (df.groupby(['A','B']).ngroup())
0    0
1    1
2    0
3   -1
4    1
5    1
6    2
dtype: int64

df['grup id'] = df.groupby(['A','B']).ngroup().replace(-1,np.nan).add(1)
print (df)
     A    B  grup id
0  1.0  2.0      1.0
1  2.0  1.0      2.0
2  1.0  2.0      1.0
3  NaN  2.0      NaN
4  2.0  1.0      2.0
5  2.0  1.0      2.0
6  2.0  2.0      3.0

Similar for replace -1 and add 1:

df['grup id'] = df.groupby(['A','B']).ngroup()
df['grup id'] = np.where(df['grup id'] == -1, np.nan, df['grup id'] + 1)
print (df)
     A    B  grup id
0  1.0  2.0      1.0
1  2.0  1.0      2.0
2  1.0  2.0      1.0
3  NaN  2.0      NaN
4  2.0  1.0      2.0
5  2.0  1.0      2.0
6  2.0  2.0      3.0

For oldiest versions of pandas (bellow 0.20.2):

df['grup id'] = df.groupby(["A","B"]).grouper.group_info[0]
df['grup id'] = np.where(df['grup id'] == -1, np.nan, df['grup id'] + 1)
print (df)
     A    B  grup id
0  1.0  2.0      1.0
1  2.0  1.0      2.0
2  1.0  2.0      1.0
3  NaN  2.0      NaN
4  2.0  1.0      2.0
5  2.0  1.0      2.0
6  2.0  2.0      3.0