根据列的总和添加新列，并在Pandas中按2个其他列分组

Question

I have the dataframe:

df = pd.DataFrame({'Continent':['North America','North America','North America','Europe','Europe','Europe','Europe'],
                'Country': ['US','Canada','Mexico','France','Germany','Spain','Italy'],
                'Status': ['Member','Non-Member','Non-Member','Member','Non-Member','Member','Non-Member'],
                'Units': [27,5,4,10,15,8,8]})

print df

       Continent  Country      Status  Units
0  North America       US      Member     27
1  North America   Canada  Non-Member      5
2  North America   Mexico  Non-Member      4
3         Europe   France      Member     10
4         Europe  Germany  Non-Member     15
5         Europe    Spain      Member      8
6         Europe    Italy  Non-Member      8

I need to add 2 columns which are summary statistics about the Continents. I need a column with the sum of Units for Member countries and Non Member countries.

so that the final output would look like:

       Continent  Member Units  Non-Member Units  Country      Status  Units
0  North America            27                 9       US      Member     27
1  North America            27                 9   Canada  Non-Member      5
2  North America            27                 9   Mexico  Non-Member      4
3         Europe            18                23   France      Member     10
4         Europe            18                23  Germany  Non-Member     15
5         Europe            18                23    Spain      Member      8
6         Europe            18                23    Italy  Non-Member      8

It seems like I need to use groupby but I can't figure out how to take the groupby values and re-insert them into the dataframe as new columns.

summary_stats = df.groupby(['Continent','Status'])['Units'].sum()
print summary_stats

Continent      Status    
Europe         Member        18
               Non-Member    23
North America  Member        27
               Non-Member     9
Name: Units, dtype: int64

I also tried not using groupby with these:

df['Member Units'] = df['Units'][df['Status'] == 'Member'].sum()
df['Non-Member Units'] = df['Units'][df['Status'] == 'Non-Member'].sum()

but that doesn't differentiate by Continent so it just adds up all the Members and Non-Members

Any help is greatly appreicated!

Answer 1

I think you need first groupby and transform sum for creating new Series all_sum . Then I think is better use numpy.where and if is member, get value from Series , if not, get 0 . Similar with non-members:

all_sum = df.groupby(['Continent','Status'])['Units'].transform(sum)
print all_sum
0    27
1     9
2     9
3    18
4    23
5    18
6    23
dtype: int64

df['Member Units'] = np.where(df['Status'] == 'Member', all_sum, 0)
df['Non-Member Units'] = np.where(df['Status'] != 'Member', all_sum, 0)
print df
       Continent  Country      Status  Units  Member Units  Non-Member Units
0  North America       US      Member     27            27                 0
1  North America   Canada  Non-Member      5             0                 9
2  North America   Mexico  Non-Member      4             0                 9
3         Europe   France      Member     10            18                 0
4         Europe  Germany  Non-Member     15             0                23
5         Europe    Spain      Member      8            18                 0
6         Europe    Italy  Non-Member      8             0                23

Answer 2

Once you have summary_stats I think you can do something like:

df['Member Units'] = summary_stats[zip(df['Continent'].values, df['Status'].values)]

The reason you need to zip the Series values is that df['Continent'] returns a series with indices, but you don't want that to happen.

Answer 3

Since you have summary_stats , you can use merge() after reshape it:

summary = summary_stats.reset_index().pivot(index='Continent', columns='Status', values='Units')

summary['Continent'] = summary.index

df = df.merge(summary, on='Continent')

Then just rename columns as you want