CODEWITHSUNDEEP
pythonlistpandas
Luis Ramon Ramirez Rodriguez
Luis Ramon Ramirez Rodriguez

Reputation: 10403

Unpack a function into a data frame

I have a function which returns two list, so a can save those in two variables like:

list_a,list_b = my_function(input)

I want to save this directly into a dataframe, something like this:

df[['list_a','list_b']]  =  my_function(input)

I got the following error:

array is not broadcastable to correct shape

Upvotes: 0

Views: 71

Answers (2)

unutbu
unutbu

Reputation: 880299

Use 

df['B'], df['C'] = my_function()

to unpack the tuple of lists returned by my_function and assign the lists to df['B'] and df['C']:

import pandas as pd
N = 5
def my_function():
    return [10]*N, [20]*N

df = pd.DataFrame({'A':[1]*N})
df['B'], df['C'] = my_function()

yields

   A   B   C
0  1  10  20
1  1  10  20
2  1  10  20
3  1  10  20
4  1  10  20

Note that the lengths of the lists returned by my_function must match the length of df.

Upvotes: 2

vk1011
vk1011

Reputation: 7179

import pandas as pd

list_a, list_b = my_function(input)
df = pd.DataFrame([list_a, list_b], columns=['a','b'])

or combined in to one line:

df = pd.DataFrame(list(my_function(input)), columns=['a','b'])

Upvotes: 0

Related Questions