如何比较两个 DataFrame 并返回包含列匹配的值的矩阵

Question

I have two data frames as follows:

df1

   id  start  end  label    item
0   1      0    3   food  burger
1   1      4    6  drink    cola
2   2      0    3   food   fries

df2 

   id  start  end  label    item
0   1      0    3   food  burger
1   1      4    6   food    cola
2   2      0    3  drink   fries

I would like to compare the two data frames (by checking where they match in the id, start, end columns) and create a matrix of size 2 x (number of items) for each id. The cells should contain the label corresponding to an item. In this example:

M_id1: [[food, drink],      M_id2: [[food], 
        [food, food]]               [drink]]

I tried looking at the pandas documentation but didn't really find anything that could help me.

Answer 1

You can merge the dataframe df1 and df2 on columns id, start, end then group the merged dataframe on id and for each group per id create key-value pairs inside dict comprehension where key is the id and value is the corresponding matrix of labels:

m = df1.merge(df2, on=['id', 'start', 'end'])
dct = {f'M_id{k}': g.filter(like='label').to_numpy().T for k, g in m.groupby('id')}

To access the matrix of labels use dictionary lookup:

>>> dct['M_id1']
array([['food', 'drink'], ['food', 'food']], dtype=object)

>>> dct['M_id2']
array([['food'], ['drink']], dtype=object)