pandas dataframe row manipulation

Question

I'm sure that I'm missing something simple, but I haven't be able to figure this one out. I have a DataFrame in Pandas with multiple rows that have the same keys, but different information. I want to place these rows onto the same row.

df = pd.DataFrame({'key': ['K0', 'K0', 'K1', 'K2'],
                  'A': ['A0', 'A1', 'A2', 'A3'],
                  'B': ['B0', 'B1', 'B2', 'B3']})

This will give me a dataframe with 4 rows and 3 columns. But there is a duplicate value 'KO' in 'key'

Is there any way to turn this into a dataframe with 3 rows, and 5 columns like shown below?

df2 = pd.DataFrame({'key': ['K0', 'K1', 'K2'],
                  'A': ['A0', 'A2', 'A3'],
                  'B': ['B0', 'B2', 'B3'],
                  'A_1': ['A1', 'NaN', 'NaN'],
                  'B_1': ['B1', 'NaN', 'NaN']})

Answer 1

Perform groupby on cumcount, then concatenate individual groups together.

gps = []
for i, g in df.groupby(df.groupby('key').cumcount()):
    gps.append(g.drop('key', 1).add_suffix(i + 1).reset_index(drop=1))

r = pd.concat(gps, 1).sort_index(axis=1)
r['key'] = df.key.unique()

r
   A1   A2  B1   B2 key
0  A0   A1  B0   B1  K0
1  A2  NaN  B2  NaN  K1
2  A3  NaN  B3  NaN  K2

You can shorten this somewhat using a list comprehension -

r = pd.concat(
         [g.drop('key', 1).add_suffix(i + 1).reset_index(drop=1) 
                    for i, g in df.groupby(df.groupby('key').cumcount())], 
         axis=1)\
      .sort_index(axis=1)

r['key'] = df.key.unique()
r
   A1   A2  B1   B2 key
0  A0   A1  B0   B1  K0
1  A2  NaN  B2  NaN  K1
2  A3  NaN  B3  NaN  K2