Slicing a DataGrameGroupBy object

user_1711146 0
Viewed 6340

Is there a way to slice a DataFrameGroupBy object?

For example, if I have:

df = pd.DataFrame({'A': [2, 1, 1, 3, 3], 'B': ['x', 'y', 'z', 'r', 'p']})

   A  B
0  2  x
1  1  y
2  1  z
3  3  r
4  3  p

dfg = df.groupby('A')

Now, the returned GroupBy object is indexed by values from A, and I would like to select a subset of it, e.g. to perform aggregation. It could be something like

dfg.loc[1:2].agg(...)

or, for a specific column,

dfg['B'].loc[1:2].agg(...)

EDIT. To make it more clear: by slicing the GroupBy object I mean accessing only a subset of groups. In the above example, the GroupBy object will contain 3 groups, for A = 1, A = 2, and A = 3. For some reasons, I may only be interested in groups for A = 1 and A = 2.

pandas-groupby pandas dataframe group-by python
2 Answers

If I understand correctly, you only want some groups, but those are supposed to be returned completely:

    A   B
1   1   y
2   1   z
0   2   x

You can solve your problem by extracting the keys and then selecting groups based on those keys.

Assuming you already know the groups:

pd.concat([dfg.get_group(1),dfg.get_group(2)])

If you don't know the group names and are just looking for random n groups, this might work:

pd.concat([dfg.get_group(n) for n in list(dict(list(dfg)).keys())[:2]])

The output in both cases is a normal DataFrame, not a DataFrameGroupBy object, so it might be smarter to first filter your DataFrame and only aggregate afterwards:

df[df['A'].isin([1,2])].groupby('A')

The same for unknown groups:

df[df['A'].isin(list(set(df['A']))[:2])].groupby('A')

I believe there are some Stackoverflow answers refering to this, like How to access pandas groupby dataframe by key

user_92944790
Related