Adding the lower levels of two Pandas MultiIndex columns

user_857390 0
Viewed 132

I have the following DataFrame:

import pandas as pd

columns = pd.MultiIndex.from_arrays([['n1', 'n1', 'n2', 'n2'],
                                     ['p',  'm',  'p',  'm']])
values = [
    [1, 2,  3,  4],
    [5, 6,  7,  8],
    [9, 10, 11, 12],
]
df = pd.DataFrame(values, columns=columns)

  n1      n2    
   p   m   p   m
0  1   2   3   4
1  5   6   7   8
2  9  10  11  12

Now I want to add another column (n3) to this DataFrame whose lower-level columns p and m should be the sums of the corresponding lower-level columns of n1 and n2:

  n1      n2      n3    
   p   m   p   m   p   m
0  1   2   3   4   4   6
1  5   6   7   8  12  14
2  9  10  11  12  20  22

Here's the code I came up with:

n3 = df[['n1', 'n2']].sum(axis=1, level=1)
level1 = df.columns.levels[1]
n3.columns = pd.MultiIndex.from_arrays([['n3'] * len(level1), level1])
df = pd.concat([df, n3], axis=1)

This does what I want, but feels very cumbersome compared to code that doesn't use MultiIndex columns:

df['n3'] = df[['n1', 'n2']].sum(axis=1)

My current code also only works for a column MultiIndex consisting of two levels, and I'd be interested in doing this for arbitrary levels.

What's a better way of doing this?

pandas
3 Answers

One way to do so with stack and unstack:

new_df = df.stack(level=1)
new_df['n3'] = new_df.sum(axis=1)
new_df.unstack(level=-1)

Output:

   n1     n2      n3    
    m  p   m   p   m   p
0   2  1   4   3   6   4
1   6  5   8   7  14  12
2  10  9  12  11  22  20
user_42384080

If you build the structure like:

df['n3','p']=1  
df['n3','m']=1

then you can write:

df['n3'] = df[['n1', 'n2']].sum(axis=1, level=1)
user_117294640

Here's another way that I just discovered which does not reorder the columns:

# Sum column-wise on level 1
s = df.loc[:, ['n1', 'n2']].sum(axis=1, level=1)

# Prepend a column level
s = pd.concat([s], keys=['n3'], axis=1)

# Add column to DataFrame
df = pd.concat([df, s], axis=1)
user_8573900
Related