Dataset from sequence of messages

Question

I have a sorted dataset by datestamp which looks like this:

user    message
A       Hi.             
B       Hello.
B       How are you?
A       I am stuck.
B       How can I help you?

What I want is to create a pandas df that would look like this:

user  message       reply
A     Hi.           Hello.
A     Hi.           How are you?
B     Hello.        I am stuck.
B     How are you?  I am stuck.
A     I am stuck.   How can I help you?

For each message, I want to find all of the replies. That means that I want the messages after current one but from the other user. How can I do this with pandas? Let's only consider a binary case of 2 users A and B.

Answer 1

First, find out when the user switch and give a separate id to each message group:

df['group_id'] = ((df['user'] != df['user'].shift()).cumsum())

user              message  group_id
   A                  Hi.         1
   B               Hello.         2
   B         How are you?         2
   A          I am stuck.         3
   B  How can I help you?         4

Then groupby each group_id and aggregate a list of the messages for each id. By shifting these messages by -1 we recive the replies for each group_id:

df_reply = df.groupby('group_id')['message'].agg(list)
df_reply = df_reply.shift(-1).reset_index().rename(columns={'message': 'reply'})

 group_id                   reply
        1  [Hello., How are you?]
        2           [I am stuck.]
        3   [How can I help you?]
        4                     NaN

The replies can then by merged back into the original dataframe. The reply lists are exploded to ensure a single reply per row:

df.merge(df_reply, on='group').explode('reply').drop('group', axis=1).dropna()

The final result:

user       message                reply
   A           Hi.               Hello.
   A           Hi.         How are you?
   B        Hello.          I am stuck.
   B  How are you?          I am stuck.
   A   I am stuck.  How can I help you?