groupby on a multiindex level and dropping the grouping level

user_4000964 0
Viewed 344

I'm doing a groupby operation on pandas dataframe, which as an example may look like this:

import pandas as pd
import numpy as np

numbers = [0, 1, 2]
colors = ['green', 'purple']
mx = pd.MultiIndex.from_product([colors, numbers], names=['color', 'numbers'])

values = np.random.rand(6)

df = pd.DataFrame(values, mx)
df = df.drop(index=("green", 2))

(index example is from pandas docs on multiindex)

Then I want to do a group operation on a level, for instance "color", but I don't want color to remain as an index in the grouped series, since it becomes redundant.

for key, series in df.groupby("color"):
    print(f"{key = }")
    print(f"Expected index: {series.index.droplevel('color').to_list()}")
    print(f"Actual index: {series.index.to_list()}")
    print()

which outputs:

key = 'green'
Expected index: [0, 1]
Actual index: [('green', 0), ('green', 1)]

key = 'purple'
Expected index: [0, 1, 2]
Actual index: [('purple', 0), ('purple', 1), ('purple', 2)]

Since the grouped level "color" is the same for every grouping, it's redundant, and I just need the level values for "numbers".

I could just drop the level inside the loop as done here, but I was wondering if there were other ways of accomplishing this that I've missed?

pandas-groupby multi-index pandas python
2 Answers

IIUC, you can simply index your slice of you dataframe with the key:

for key, series in df.groupby("color"):
    print(f'Key: {key}')
    print(series.loc[key])
    print('\n')

Output:

Key: green
                0
numbers          
0        0.913962
1        0.457205


Key: purple
                0
numbers          
0        0.939128
1        0.778389
2        0.715971
user_63615310

It looks like you want to loop through first level value. You can just try:

for color in df.index.get_level_values('color').unique():
    series = df.loc[color]
    print(color)
    print(series)

Output:

green
                0
numbers          
0        0.161444
1        0.403312
purple
                0
numbers          
0        0.369750
1        0.222223
2        0.565426
user_42384080
Related