Is it possible to pivot a pandas dataframe with a column of lists?

Question

I have a dataframe with a column of lists that I need to reshape from long to wide.

df_have = pd.DataFrame({'id': [1, 1, 2, 2, 3, 3],
                        'year': [[2009], [2010], [2012, 2014], [2015], [2016, 2017], [2018]]})
df_have

    id  year
0   1   [2009]
1   1   [2010]
2   2   [2012, 2014]
3   2   [2015]
4   3   [2016, 2017]
5   3   [2018]

Using DataFrame.pivot as in:

df_have.pivot(index='id', columns='year')

gives this error: TypeError: unhashable type: 'list'

How can I reshape the dataframe into this format?

df_want = pd.DataFrame({'id': [1, 2, 3],
                        'year_1': [[2009], [2012, 2014], [2016, 2017]],
                        'year_2': [[2010], [2015], [2018]]})
df_want

    id  year_1          year_2
0   1   [2009]          [2010]
1   2   [2012, 2014]    [2015]
2   3   [2016, 2017]    [2018]

`df.assign(col=df.groupby('id').cumcount()).pivot(index='id', columns='col', values='year')`. — Quang Hoang, Jan 29 '20 at 20:33

score 0 · Accepted Answer · answered Jan 29 '20 at 20:41

unstack would have been easier:

# For each `id` give the row a `year_1`, `year_2`, etc.
df_have['year_name'] = df_have.groupby('id').transform(lambda group: [f'year_{i}' for i in range(1, group.shape[0] + 1)])

df_have.set_index(['id', 'year_name']).unstack()

Result

                           year        
year_name        year_1  year_2
id                             
1                [2009]  [2010]
2          [2012, 2014]  [2015]
3          [2016, 2017]  [2018]

Is it possible to pivot a pandas dataframe with a column of lists?

1 Answers1