Merge two datasets that have lists and keep the list after merge using pandas

Viewed 60

I have this two dataframes that I am having a hard time to merge:

    df1 = pd.DataFrame({'id': [ ["001",  "001"], ["001"], ["007",   "001"]]})

OUTPUT:

    id
    0   [001, 001]
    1   [001]
    2   [007, 001]

and

df2 = pd.DataFrame({'id': [ "001", "007"],'name': ['Name01', 'Name02']})

OUTPUT:

id  name
0   001 Name01
1   007 Name02

And I want to arrive is this:

df3 = pd.DataFrame({'id':  [ ["001",  "001"], ["001"], ["007",   "01"]],
                    'name': [ ['Name01','Name01'], ['Name01'], ['Name02', 'Name01']]})

OUTPUT:

    id  name
0   [001, 001]  [Name01, Name01]
1   [001]   [Name01]
2   [007, 01]   [Name02, Name01]

My problem is that I can merge, but I am not managing to put in the format I want. What I have at the moment is here:

pd.DataFrame(df2.merge(df1.explode('id'), on= 'id')).groupby('id').agg(lambda x: x.tolist())

OUTPUT:

id   name
001 [Name01, Name01, Name01, Name01]
007 [Name02]
2 Answers

Use mapping by dictionary created by df2 in list comprehension, what should be faster like explode and aggregate list, the best test in real data:

d = df2.set_index('id')['name'].to_dict()
df1['name'] = [[d[y] for y in x if y in d] for x in df1['id']]
print (df1)
           id              name
0  [001, 001]  [Name01, Name01]
1       [001]          [Name01]
2  [007, 001]  [Name02, Name01]

We can do explode + merge

df1=df1.explode('id').reset_index().merge(df2,how='left').groupby('index').agg(list)
               id              name
index                              
0      [001, 001]  [Name01, Name01]
1           [001]          [Name01]
2      [007, 001]  [Name02, Name01]

Or just map and assign

df1['name']=df1.id.explode().map(df2.set_index('id').name).groupby(level=0).agg(list)
0    [Name01, Name01]
1            [Name01]
2    [Name02, Name01]
Name: id, dtype: object
Related