CODEWITHSUNDEEP
pythonpandaspandas-groupbymulti-index
sunny
sunny

Reputation: 748

Pandas groupby rows into list and sum

I have a multi-index pandas dataframe with the below data:

          F     M   
         0  5  10 30
x  y  a1  0  1  0  1
      a2  1  0  0  0
      a3  0  1  0  1
      a4  0  1  0  2
x1 y1 a1  0  1  1  4
      a2  0  1  0  1
      a3  1  1  0  1
      a4  2  1  1  2
x2 y2 a1  0  0  0  0
      a2  0  0  0  0

I want to group the data through (A, B) and consolidate column C in a list along with the sum of the rows.

Required output:

              sex             F               M
              counts          0  5  10  30    0  5  10  30
     A  B  C
     x  y  [a1,a2,a3, a4]     1  3   3   1    0  3   0   4

     x1 y1 [a1, a2, a3, a4]   3  4   4   0    3  0   1   8

     ......

I have tried df.groupby(['A', 'B']).sum() but it ignores column C in the output dataframe.

Code To Reproduce 

index = pd.MultiIndex(levels=[['x', 'x1', 'x2'], ['y', 'y1', 'y2'], ['a1', 'a2', 'a3', 'a4']],
           labels=[[0, 0, 0, 0, 1, 1, 1, 1, 2, 2], [0, 0, 0, 0, 1, 1, 1, 1, 2, 2], [0, 1, 2, 3, 0, 1, 2, 3, 0, 1]])

columns = pd.MultiIndex(levels=[['F', 'M'], [0, 5, 10, 30]],
           labels=[[0, 0, 1, 1], [0, 1, 2, 3]])

data = np.array([
       [0, 1, 0, 1],
       [1, 0, 0, 0],
       [0, 1, 0, 1],
       [0, 1, 0, 2],
       [0, 1, 1, 4],
       [0, 1, 0, 1],
       [1, 1, 0, 1],
       [2, 1, 1, 2],
       [0, 0, 0, 0],
       [0, 0, 0, 0]])

df = pd.DataFrame(data, index=index, columns=columns)

Upvotes: 1

Views: 506

Answers (1)

cs95
cs95

Reputation: 402263

Use groupby and sum for aggregating the numeric data, and apply(tuple) for aggregating the index level.

g = df.reset_index(level=-1).groupby(level=[0, 1])
res = g.sum().set_index(g.level_2.apply(tuple), append=True)


print(res)
                          F     M   
                       0  5  10 30
      level_2                     
x  y  (a1, a2, a3, a4)  1  3  0  4
x1 y1 (a1, a2, a3, a4)  3  4  2  8
x2 y2 (a1, a2)          0  0  0  0

Note, the index can only contain hashable values, and lists are not hashable, so tuples are the next best thing.

Upvotes: 2

Related Questions