Панды выбирают первые x строк, соответствующих значениям y, удаляя результаты ниже x

Question

У меня есть такой фрейм данных:

ID A B
0  7 4
0  5 2
0  0 3
1  6 7
1  8 9
2  5 5

Я хотел бы выбрать первые x строк для всех IDs, но только с тем, что для этих IDs больше строк, например:

Если x == 2:

ID A B
0  7 4
0  5 2
1  6 7
1  8 9

Если x == 3:

ID A B
0  7 4
0  5 2
0  0 3

... и т. Д.

Использованиеdf.groupby("ID").head(2) приблизительно соответствует тому, что я хочу, но включает в себя первую строку для ID "2", которая мне не нужна:

ID A B
0  7 4
0  5 2
1  6 7
1  8 9
2  5 5

Существует ли эффективный способ сделать это, не прибегая к подсчету строкза каждый ID?

Answer 1

Используйте следующий код:

x = 2
gr = df.groupby('ID', as_index=False)\
    .apply(lambda grp: grp.head(x) if len(grp) >= x else None)\
    .reset_index(drop=True)

Применяемая здесь лямбда-функция проверяет, соответствует ли длина группы не менее x (вид фильтрации по длине группы) и для таких групп выводит первые x строк.

Таким образом, вы избегаете второго groupby.

Результат:

   ID  A  B
0   0  7  4
1   0  5  2
2   1  6  7
3   1  8  9

Answer 2

Использование groupby + duplicated с keep=False:

v = df.groupby('ID').head(2)
v[v.ID.duplicated(keep=False)]

   ID  A  B
0   0  7  4
1   0  5  2
3   1  6  7
4   1  8  9

---

Вы также можете сделать 2x groupby (нет ... не рекомендую):

df[df.groupby('ID').ID.transform('size').gt(1)].groupby('ID').head(2)

   ID  A  B
0   0  7  4
1   0  5  2
3   1  6  7
4   1  8  9