панды - количество элементов после группового на нескольких столбцов

Question

У меня есть данные фрейма данных, сгруппированные по двум столбцам (X, Y), и затем у меня есть количество элементов в Z. Идея здесь состоит в том, чтобы найти 2 верхних числа элементов в X, Y.

Фрейм данных долженвыглядеть так:

mostCountYInX = df.groupby(['X','Y'],as_index=False).count()
 C      X   Y     Z
USA    NY  NY    5
USA    NY  BR    14
USA    NJ  JC    40
USA    FL  MI    3
IND    MAH MUM   4
IND    KAR BLR   2
IND    KER TVM   2
CHN    HK  HK    3
CHN    SH  SH    3

По отдельности я могу извлечь информацию, которую ищу:

XTopCountInTopY = mostCountYInX[mostCountYInX['X'] == 'NY']
XTopCountInTopY = XTopCountInTopY.nlargest(2,'Y')

В группе, которую я искал, я знал, что X = NY, и получил2 верхние записи.Есть ли способ напечатать их вместе?

Скажите, что меня интересуют ИНД и США, тогда ожидаемый результат:

 C     X   Y     Z
USA    NJ  JC    40
USA    NY  BR    14
IND    MAH MUM   4
IND    KAR BLR   2

Answer 1

Я думаю, вам нужно groupby по индексу с параметром sort=False, затем apply с использованием lambda function и sort_values по Z с использованием параметра ascending=Falseзатем возьмите 2 верхних значения и reset_index как:

mask = df.index.isin(['USA','IND'])
df = df[mask].groupby(df[mask].index,sort=False).\
              apply(lambda x: x.sort_values('Z',ascending=False)[:2]).\
              reset_index(level=0,drop=True)
print(df)

       X    Y   Z
USA   NJ   JC  40
USA   NY   BR  14
IND  MAH  MUM   4
IND  KAR  BLR   2

РЕДАКТИРОВАТЬ: После того, как OP изменил кадр данных:

mask = df['C'].isin(['USA','IND'])
df = df[mask].groupby('C',sort=False).\
              apply(lambda x: x.sort_values('Z',ascending=False)[:2]).\
              reset_index(drop=True)

print(df)
     C    X    Y   Z
0  USA   NJ   JC  40
1  USA   NY   BR  14
2  IND  MAH  MUM   4
3  IND  KAR  BLR   2