CODEWITHSUNDEEP
pythonpandaslistdataframe
jason
jason

Reputation: 2106

How to merge each row of dataframe into a list in python

I have a dataframe in python

import pandas as pd
d = {'name':['a','b','c','d','e'],'location1': [1, 2,3,8,6], 'location2': 
[2,1,4,6,8]}
df = pd.DataFrame(data=d)

df is as follow:

    name   location1  location2
0    a          1          2
1    b          2          1
2    c          3          4
3    d          8          6
4    e          6          8

I try to obtain a dataframe as:

    name    loc
0    a     [1, 2]
1    b     [2, 1]
2    c     [3, 4]
3    d     [8, 6]
4    e     [6, 8]

How to efficiently convert that?

Upvotes: 1

Views: 53

Answers (1)

cs95
cs95

Reputation: 402363

Here are some suggestions.

Listification and Assignment

# pandas >= 0.24
df['loc'] = df[['location1', 'location2']].to_numpy().tolist()
# pandas < 0.24
df['loc'] = df[['location1', 'location2']].values.tolist()
df

  name  location1  location2     loc
0    a          1          2  [1, 2]
1    b          2          1  [2, 1]
2    c          3          4  [3, 4]
3    d          8          6  [8, 6]
4    e          6          8  [6, 8]

Remove the columns using drop.

(df.drop(['location1', 'location2'], 1)
   .assign(loc=df[['location1', 'location2']].to_numpy().tolist()))

  name     loc
0    a  [1, 2]
1    b  [2, 1]
2    c  [3, 4]
3    d  [8, 6]
4    e  [6, 8]

zip with pop using List Comprehension

df['loc'] = [[x, y] for x, y in zip(df.pop('location1'), df.pop('location2'))]
# or 
df['loc'] = [*map(list, zip(df.pop('location1'), df.pop('location2')))]
df
  name     loc
0    a  [1, 2]
1    b  [2, 1]
2    c  [3, 4]
3    d  [8, 6]
4    e  [6, 8]

pop destructively removes the columns, so you get to assign and cleanup in a single step.

Upvotes: 1

Related Questions