Фильтр групп по тому, равно ли максимальное значение одного столбца, с учетом условия другого, определенному значению

Question

Я хочу нарезать (или отфильтровать) те группы (хотя я бы согласился на строки.), Которые принадлежат группам в d2.groupby('g'), где максимум значений s в группе, для строк, где test равно 1, равно 2.

pd.__version__
# '0.24.1'

d2 = pd.DataFrame({'g': [1, 1, 1, 2, 2, 2],
                   'test': [1, 1, 0, 1, 0, 0],
                   's': [1, 2, 3, 1, 2, 3]})

d2
   g  test  s
0  1     1  1
1  1     1  2
2  1     0  3
3  2     1  1
4  2     0  2
5  2     0  3

Простым английским языком: Я хочу сохранить группы, в которых максимальное значение s с соответствующим test значением 1 равно 2. Для приведенного ниже примера я хочу сохранить g группа 1, поскольку во втором ряду есть test == 1 и s == 2, а в третьем ряду - s == 0. Ожидаемый результат для этих данных:

   g  test  s
0  1     1  1
1  1     1  2
2  1     0  3

Я пробовал d2.groupby('g').filter(lambda x: (x.test == 1)) # followed by nonsense, а также d2.groupby('g')[[d2.s == 1]] # with more nonsense. Последний дает мне сообщение о том, что объекты Series являются изменяемыми и не подлежащими изменению. Я пробовал ряд других, одинаково бессмысленных и бесполезных подходов. Как мне groupby('g') и отфильтровать по максимальному значению s, где test равно 1? Кажется, это должно быть легко, но я не могу его получить. Я мог бы добраться туда, добавив столбец, но это похоже на кучу.

Примечание: я искал ответ для этого. Я бы с удовольствием воспользовался любыми подсказками по поиску того, как найти текущие решения для этого вопроса, если таковые имеются. Спасибо!

Answer 1

Использование filter

d2.groupby('g').filter(lambda x : (x['s']*x['test']).max()==2)
Out[390]: 
   g  s  test
0  1  1     1
1  1  2     1
2  1  3     0

---

Другой способ использования isin , поскольку filter обычно медленнее в большом кадре данных

s=d2.s.mul(d2.test).groupby(d2['g']).max()==2
d2.loc[d2.g.isin(s.index[s])]
Out[394]: 
   g  s  test
0  1  1     1
1  1  2     1
2  1  3     0

Answer 2

Вы можете использовать groupby и transform для вычисления маски, что-то вроде этого:

df[df['s'].where(df['test'] == 1).groupby(df['g']).transform('max') == 2]

   g  test  s
0  1     1  1
1  1     1  2
2  1     0  3

Идея Series.where состоит в том, чтобы позволить нам рассматривать только значения "s", где "test" равен 1.

Вот аналогичная версия выше, вдохновленная WeNYoBen, которая будет работать, потому что здесь «test» равен 0 или 1.

df[df['s'].mul(df['test']).groupby(df['g']).transform('max').eq(2)]

   g  test  s
0  1     1  1
1  1     1  2
2  1     0  3