Any straightforward way of adding multiindex columns?

Viewed 61
df = pd.DataFrame(data=np.array([[80, 82, 81, 84], [100, 101, 98, 102]]).T, 
                  index=[0, 1, 2, 3], 
                  columns=pd.MultiIndex.from_product([['weight'], ['Apple', 'Orange']]))

  weight       
   Apple Orange
0     80    100
1     82    101
2     81     98
3     84    102

I would like to put their .diff() into a new column, but I cannot achieve that naturally like working with a simple df with single-level columns:

# 1st try:
df['diff'] = df['weight'].diff()

KeyError: 'diff'
ValueError: Wrong number of items passed 2, placement implies 1
# 2nd try:
df[['diff']] = df[['weight']].diff()

ValueError: Columns must be same length as key
# 3rd try:
df = df.join(df[['weight']].diff().rename(columns={'weight':'diff'}))

  weight         diff       
   Apple Orange Apple Orange
0     80    100   NaN    NaN
1     82    101   2.0    1.0
2     81     98  -1.0   -3.0
3     84    102   3.0    4.0

The 3rd try works but it seems so unnatural to me. I assume there is a more straightforward way to add multiindex columns without doing renaming columns and join, merge or concat?

Thanks a lot!

2 Answers

You can use MultiIndex.set_levels:

In [3301]: x = df.diff()

In [3321]: x.columns = x.columns.set_levels(['diff'], level=0)

In [3324]: pd.concat([df, x], 1)
Out[3324]: 
  weight         diff       
   Apple Orange Apple Orange
0     80    100   NaN    NaN
1     82    101   2.0    1.0
2     81     98  -1.0   -3.0
3     84    102   3.0    4.0

df.diff() returns a dataframe, if it would return a Series then you could easily assign as a column, otherwise you need to .join() or .concat()

Related