Python Dataframe add two columns containing lists

Question

I have a dataframe of two columns, each containing list as elements. I want to perform rowwise element additon of two lists. My below solution is inspired from this answer Element-wise addition of 2 lists?

My code:

df = pd.DataFrame({'A':[[1,2,3],[10,20]],'B':[[4,5,6],[40,50]]})

           A          B
0  [1, 2, 3]  [4, 5, 6]
1   [10, 20]   [40, 50]

df['C'] = df[['A','B']].apply(lambda x: list(map(add,x[0],x[1])))
df['C'] = 
    A   B
0  11  44
1  22  55

Expected answer:

df['C'] = 
   C
0 [5,7,9]
1 [50,70]

score 1 · Accepted Answer · 2022-04-30 20:32:11Z

1

Since you need to apply the function row-wise, you just need axis=1:

from operator import add
df['C'] = df[['A','B']].apply(lambda x: list(map(add,x[0],x[1])), axis=1)

or

df['C'] = df[['A','B']].apply(lambda x: [i+j for i,j in zip(*x)], axis=1)

An alternative option is to explode the lists; sum; then groupby.agg to get the lists back in:

df['C'] = df.explode(['A','B']).sum(axis=1).astype(int).groupby(level=0).agg(list)

Output:

           A          B          C
0  [1, 2, 3]  [4, 5, 6]  [5, 7, 9]
1   [10, 20]   [40, 50]   [50, 70]

edited Apr 30, 2022 at 20:32

answered Apr 30, 2022 at 6:59

user7864386

Sign up to request clarification or add additional context in comments.

1 Comment

Mainland Over a year ago

It worked. Thanks. I am happy to have come this close to the actual answer. Also, I found an alternative answer as df[['A','B']].apply(lambda x: [sum(x) for x in zip(x[0],x[1])],axis=1)

Collectives™ on Stack Overflow

Python Dataframe add two columns containing lists

1 Answer 1

1 Comment

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Linked

Related