python: извлекать строки, значение столбца которых больше 3

Question

существует следующий фрейм данных:

id  year number
1   2016  3
1   2017  5
2   2016  1
2   2017  5
...

Я хочу извлечь строки с идентификатором groupby и значением столбца числа больше 3 как в 2016, так и в 2017 году.

например, в приведенных выше первых 4 строках результат:

   id  year number
    1   2016  3
    1   2017  5

Спасибо!

Answer 1

Сравните на >=3 и используйте GroupBy.transform для серии с таким же размером, как у оригинала, поэтому возможный фильтр по boolean indexing:

df1 = df[(df["number"] >= 3).groupby(df["id"]).transform('all')]
#alternative for reassign mask to column
#df = df[df.assign(number= df["number"] >= 3).groupby("id")['number'].transform('all')]
print (df1)
   id  year  number
0   1  2016       3
1   1  2017       5

Или используйте filter, но это должно быть медленно, если большой DataFrame или много групп:

df1 = df.groupby("id").filter(lambda x: (x["number"] >= 3).all())

Answer 2

>>> great_in_both_years = df.groupby("id").apply(lambda x: (x["number"] >= 3).all())
>>> great_in_both_years
id
1     True
2    False
dtype: bool

>>> df.loc[lambda x: x["id"].map(great_in_both_years)]
   id  year  number
0   1  2016       3
1   1  2017       5