Условие в строках, содержащихся в группах

Question

             A      B   C
0   2002-01-12  Sarah  39
1   2002-01-12   John  17
2   2002-01-12  Susan  30
3   2002-01-15  Danny  12
4   2002-01-15  Peter  25
5   2002-01-15   John  25
6   2002-01-20   John  16
7   2002-01-20   Hung  10
8   2002-02-20   John  20
9   2002-02-20  Susan  40
10  2002-02-24  Rebel  40
11  2002-02-24  Susan  15
12  2002-02-24   Mark  38
13  2002-02-24  Susan  30

Я хочу выбрать полные A группы, которые содержат John и Susan.

Вывод должен быть:

             A      B   C
0   2002-01-12  Sarah  39
1   2002-01-12   John  17
2   2002-01-12  Susan  30
6   2002-01-20   John  16
7   2002-01-20   Hung  10
8   2002-02-20   John  20
9   2002-02-20  Susan  40

Я пробовал:

df.groupby('A').apply(lambda x: ((df.B == x.John) & (df.B == x.Susan)))

Answer 1

создайте массив дат как пересечение дат, которые содержат John и даты, которые содержат Susan:

dates = np.intersect1d(
    df.A.values[df.B.values == 'John'], 
    df.A.values[df.B.values == 'Susan']
)

, затем используйте массив дат для фильтрации кадра данных

df[df.A.isin(dates)]

# outputs:
            A      B   C
0  2002-01-12  Sarah  39
1  2002-01-12   John  17
2  2002-01-12  Susan  30
8  2002-02-20   John  20
9  2002-02-20  Susan  40

Время:

Сравнение решений, предоставленных jpp , ALollz и выше:

Решение на основе numpy в несколько раз эффективнеечем другие.

In [288]: def hal(df):
     ...:     dates = np.intersect1d(
     ...:      df.A.values[df.B.values == 'John'], 
     ...:      df.A.values[df.B.values == 'Susan']
     ...:     )
     ...:     return df[df.A.isin(dates)]
     ...:

In [289]: def jpp(df):
     ...:     s = df.groupby('A')['B'].apply(set)
     ...:     return df[df['A'].map(s) >= {'John', 'Susan'}]
     ...:

In [290]: def alollz(df):
     ...:     flag = df.groupby('A').B.transform(lambda x: ((x=='Susan').any() & (x == 'John').any()).sum().astype('boo
     ...: l'))
     ...:     return df[flag==True]
     ...:

In [291]: %timeit hal(df)
394 µs ± 6.42 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [292]: %timeit jpp(df)
1.46 ms ± 27.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [293]: %timeit alollz(df)
4.9 ms ± 75 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Однако решение, предложенное ALollz, можно ускорить в 2 раза, исключив некоторые лишние ненужные операции и сократившись до пустых массивов для сравнения.

In [294]: def alollz_improved(df):
     ...:     v = df.groupby('A').B.transform(lambda x: (x.values=='Susan').any() & (x.values=='John').any())
     ...:     return df[v]
     ...:

In [295]: %timeit alollz_improved(df)
2.2 ms ± 38.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Answer 2

Создать серию, сопоставляющую каждую дату с set именами.Затем используйте set.issuperset через синтаксический сахар >=:

s = df.groupby('A')['B'].apply(set)

res = df[df['A'].map(s) >= {'John', 'Susan'}]

print(res)

            A      B   C
0  2002-01-12  Sarah  39
1  2002-01-12   John  17
2  2002-01-12  Susan  30
8  2002-02-20   John  20
9  2002-02-20  Susan  40

Answer 3

Вы можете использовать groupby + transform, чтобы создать флаг для групп, которые удовлетворяют этому условию.И тогда вы можете замаскировать оригинал df под этим флагом.Если вы не хотите изменять исходный df, вы можете создать отдельный Series с именем flag, в противном случае вы также можете просто назначить его столбцу в исходном df

import pandas as pd
# As Haleemur Ali points out, use x.values to make it faster
flag = df.groupby('A').B.transform(lambda x: (x.values == 'Susan').any() & (x.values == 'John').any())

Тогда вы можете отфильтровать df

df[flag]
#            A      B   C
#0  2002-01-12  Sarah  39
#1  2002-01-12   John  17
#2  2002-01-12  Susan  30
#8  2002-02-20   John  20
#9  2002-02-20  Susan  40