Brian H. Truong
Brian H. Truong

Reputation: 53

Python pandas How to perform this kind of pivot?

For example, say I have a dataframe (id is letter):

letter|color |number
a     |green |2
a     |blue  |3
b     |red   |3
b     |blue  |4
b     |yellow|1
c     |red   |9
c     |blue  |5

What I want is to transform it to:

letter|color_1|color_2|color_3|number_1|number_2|number_3
a     |green  |blue   |       |2       |3       |
b     |red    |blue   |yellow |3       |4       |1
c     |red    |blue   |       |9       |5       |

Upvotes: 0

Views: 42

Answers (2)

user3483203
user3483203

Reputation: 51165

You can stack before you pivot, with some data manipulation:

s = df.set_index('letter').stack().reset_index(name='vals')
counter = s.groupby(['letter', 'level_1']).cumcount().add(1).astype(str)

(s.assign(flag=s.level_1 + '_' + counter)
    .pivot_table(index='letter', values='vals', columns='flag', aggfunc='first'))

tmp    color_1 color_2 color_3 number_1 number_2 number_3
letter
a        green    blue     NaN        2        3      NaN
b          red    blue  yellow        3        4        1
c          red    blue     NaN        9        5      NaN

Upvotes: 1

jpp
jpp

Reputation: 164773

You can perform a GroupBy + agg operation and create a new dataframe concatenating the results:

g = df.groupby('letter').agg(list)

df1 = pd.DataFrame(g['color'].values.tolist(), index=g.index).add_prefix('color_')
df2 = pd.DataFrame(g['number'].values.tolist(), index=g.index).add_prefix('number_')

res = pd.concat([df1, df2], axis=1).reset_index()

print(res)

   letter color_0 color_1 color_2  number_0  number_1  number_2
0  a        green    blue    None         2         3       NaN
1  b          red    blue  yellow         3         4       1.0
2  c          red    blue    None         9         5       NaN

Upvotes: 1

Related Questions