user10250831
user10250831

Reputation: 25

Add column using groupby in multiindex Pandas

I am trying to find the sum of a column based upon the groupby function. so in this example i want to find the sum of all the bar, baz, foo, and qux.

the sum would be added to a new column at the end. i can get the results i need but i can not join it back to the dataframe.

import numpy as np
import pandas as pd


arrays = [np.array(['bar', 'bar', 'baz', 'baz', 'foo', 'foo', 'qux', 'qux']),
          np.array(['one', 'two', 'one', 'two', 'one', 'two', 'one', 'two'])]

np.random.seed(7) 
df = pd.DataFrame(np.random.randn(8, 4), index=arrays)

results=df.groupby(level=[0]).sum(axis=1)

col_names=results.columns.values
hold=[]
for i in col_names:
    hold.append('sum_'+str(i))
results.columns=hold

df=pd.concat([df,results],axis=1)

Desired result below. thanks for looking

        0          1       2       3    sum_0   sum_1   sum_2   sum_3
bar one  1.69    (0.47)  0.03    0.41    0.90    (0.46)  0.03    (1.35)
bar two  (0.79)  0.00    (0.00)  (1.75)  0.90    (0.46)  0.03    (1.35)
baz one  1.02    0.60    (0.63)  (0.17)  1.52    0.34    (0.87)  (1.62)
baz two  0.51    (0.26)  (0.24)  (1.45)  1.52    0.34    (0.87)  (1.62)
foo one  0.55    0.12    0.27    (1.53)  2.21    0.28    (0.11)  0.50 
foo two  1.65    0.15    (0.39)  2.03    2.21    0.28    (0.11)  0.50 
qux one  (0.05)  (1.45)  (0.41)  (2.29)  1.00    (1.87)  (1.15)  (1.22)
qux two  1.05    (0.42)  (0.74)  1.07    1.00    (1.87)  (1.15)  (1.22)

Upvotes: 1

Views: 233

Answers (1)

cs95
cs95

Reputation: 402844

Use transform instead, you can rid your code of that loop.

df = pd.concat([df, df.groupby(level=0).transform('sum').add_prefix('sum_')], axis=1)
df
             0      1      2      3 sum_0  sum_1  sum_2  sum_3
bar one   1.69  -0.47   0.03   0.41  0.90  -0.46   0.03  -1.35
    two  -0.79   0.00  -0.00  -1.75  0.90  -0.46   0.03  -1.35
baz one   1.02   0.60  -0.63  -0.17  1.52   0.34  -0.87  -1.62
    two   0.51  -0.26  -0.24  -1.45  1.52   0.34  -0.87  -1.62
foo one   0.55   0.12   0.27  -1.53  2.21   0.28  -0.11   0.50
    two   1.65   0.15  -0.39   2.03  2.21   0.28  -0.11   0.50
qux one  -0.05  -1.45  -0.41  -2.29  1.00  -1.87  -1.15  -1.22
    two   1.05  -0.42  -0.74   1.07  1.00  -1.87  -1.15  -1.22

Upvotes: 1

Related Questions