Label sequences by group in Pandas

user_8103192 0
Viewed 185

I have the following dataframe :

df = pd.DataFrame({'ID_1' : ['i1','i1','i1','i1','i1','i1','i1','i1'],
             'ID_2' : ['a1','a1','a1','a1','a2','a2','a2','a2'],
             'label':['a','b','b','a','a','a','a','b']})

I would like to get an id for the sequences of 'label' for each group ('ID_1' , 'ID_2'), so that the output would be :

ID_1    ID_2    label   seq_id
i1      a1      a       1
i1      a1      b       2
i1      a1      b       2
i1      a1      a       3
i1      a2      a       1
i1      a2      a       1
i1      a2      a       1
i1      a2      b       2

If the group was not a concern, this would have worked :

df['id_seq'] = df['label'].shift(1) != df['label']).cumsum()

So I tried to adapt it with the following :

df['shifted_label'] = df.groupby(['ID_1','ID_2']).label.shift(1)
def get_sequence_group(grouped) :
    return (grouped['shifted_label'] != grouped['label']).cumsum()
df.groupby(['ID_1','ID_2']).apply(get_sequence_group)

But it's not working at all. Would like to understand why and how to proceed. Thanks !

pandas-groupby pandas dataframe python
2 Answers

Try this. I think you were close; use groupby and transform with your logic:

df['seq_id'] = df.groupby(['ID_1', 'ID_2'])['label']\
                 .transform(lambda x: (x != x.shift()).cumsum())

Output:

  ID_1 ID_2 label  seq_id
0   i1   a1     a       1
1   i1   a1     b       2
2   i1   a1     b       2
3   i1   a1     a       3
4   i1   a2     a       1
5   i1   a2     a       1
6   i1   a2     a       1
7   i1   a2     b       2
user_63615310

To keep it completely vectorized, we can groupby on the diff with shift. This way we avoid using apply

diff = df['label'].ne(df['label'].shift())
df['seq_id'] = diff.groupby([df['ID_1'], df['ID_2']]).cumsum()

  ID_1 ID_2 label  seq_id
0   i1   a1     a       1
1   i1   a1     b       2
2   i1   a1     b       2
3   i1   a1     a       3
4   i1   a2     a       0
5   i1   a2     a       0
6   i1   a2     a       0
7   i1   a2     b       1
user_90812670
Related