How to flatten rows having similar index in pandas?

user_4800652 0
Viewed 208

I have a dataframe like 

df = pd.DataFrame({'a':[np.array([5,6]),6,np.array([8,10]),7],'b':[np.array([7,8]),9,np.array([15,10]),7]},index=[0,0,1,1])

        a         b
0   [5, 6]    [7, 8]
0        6         9
1  [8, 10]  [15, 10]
1        7         7

When I try groupby

df.groupby(level=0).apply(lambda x: pd.Series(x.values.flatten()))

         0         1  2  3
0   [5, 6]    [7, 8]  6  9
1  [8, 10]  [15, 10]  7  7

So how to use apply in such a way that I end up flattening the cells with similar index under the same column. 

       a         b
0   [5, 6,6]    [7, 8,9]
1  [8, 10,7]  [15, 10,7]
pandas numpy python
1 Answers

This is a job for numpy.hstack. However, getting the output of a groupby into a dataframe is always a bit tricky when the values are multidimensional. Fitting things into a series usually works:

df.groupby(level=0).apply(lambda g: pd.Series({
    'a': np.hstack(g['a'].values), 
    'b': np.hstack(g['b'].values)
}))

Of course, enumerating the dictionary would be nicer...

For n columns a dict comprehension would be better i.e 

df.groupby(level=0).apply(lambda g: pd.Series({i: np.hstack(g[i].values) for i in df.columns}))
user_52767970
Related