pandas：在每组中按复杂条件选择子数据集答案

【问题标题】：pandas: select sub-datasets by complex condition in each grouppandas：在每组中按复杂条件选择子数据集
【发布时间】：2021-03-11 12:55:02
【问题描述】：

我需要从给定的数据框中选择一个子集。这是df：

import pandas as pd
import numpy as np

df = pd.DataFrame({
                   'custom_id': ['aa','aa','aa','aa','aa','aa',
                                 'bk', 'bk', 'bk', 'bk','bk',
                                 'dd', 'dd', 'dd', 'dd', 'dd',
                                 'ff', 'ff', 'ff', 'ff', 'ff', 'ff',
                                 'pu', 'pu', 'pu', 'pu'],
                   'sending_num': [11, 252, 198, 266, 5317, 'from',
                                   67, 287, 909, 881, 'from',
                                   22, 55, 'from', 376, 98,
                                   901, 126, 22, 381, 867, 'from',
                                   421, 81, 326, 'from'],
                   'receiving_num': [900, 11, 252, 198, 266, 5317,
                                     345, 67, 287, 909, 881,
                                     432, 22, 55, 65, 376,
                                     42, 901, 126, 22, 381, 867,
                                     66, 421, 81, 326],
                   'note': [np.nan, 'flag', np.nan, np.nan, 'flag', np.nan,
                            'flag', np.nan, np.nan, np.nan, np.nan,
                            np.nan, 'flag', np.nan, np.nan, np.nan,
                            np.nan, np.nan, np.nan, np.nan, 'flag', np.nan,
                            np.nan, np.nan, np.nan, np.nan]
                   })

df 是这样的：

   custom_id sending_num  receiving_num  note
0         aa          11            900   NaN
1         aa         252             11  flag
2         aa         198            252   NaN
3         aa         266            198   NaN
4         aa        5317            266  **flag**
5         aa        **from**        5317   NaN
6         bk          67            345  flag
7         bk         287             67   NaN
8         bk         909            287   NaN
9         bk         881            909   NaN
10        bk        from            881   NaN
11        dd          22            432   NaN
12        dd          55             22  **flag**
13        dd        **from**         55   NaN
14        dd         376             65   NaN
15        dd          98            376   NaN
16        ff         901             42   NaN
17        ff         126            901   NaN
18        ff          22            126   NaN
19        ff         381             22   NaN
20        ff         867            381  **flag**
21        ff        **from**        867   NaN
22        pu         421             66   NaN
23        pu          81            421   NaN
24        pu         326             81   NaN
25        pu        from            326   NaN

我希望根据以下规则选择一个子集：对于每个组（自定义 id），如果：行中出现'from'，并且在其'note'列中，上面的行有'标志的价值。例如，对于组 'aa'，在其 'sending_num' 列中有一个 'from'，同时，在其上面的行（第 4 行）中，在同一组的 'note' 列中有一个 'flag' ，所以 'aa' 是一个目标；类似于组 'dd' 和 'ff' 因为在它们的 'sending_num' 列中有 'from'，而在 'note' 列的上一行中，有 'flag'，所以这两个被选中，但不是其他组。我尝试编写一个循环和 iloc 来执行此操作，但速度很慢。最终，我希望根据规则有一个这样的子集：

   custom_id sending_num  receiving_num  note
0         aa          11            900   NaN
1         aa         252             11  flag
2         aa         198            252   NaN
3         aa         266            198   NaN
4         aa        5317            266  flag # 'flag' row &
5         aa        from           5317   NaN # 'from' row are adjacent for 'aa'
6         dd          22            432   NaN
7         dd          55             22  flag # 'flag' row &
8         dd        from             55   NaN # 'from' row are adjacent for 'dd'
9         dd         376             65   NaN
10        dd          98            376   NaN
11        ff         901             42   NaN
12        ff         126            901   NaN
13        ff          22            126   NaN
14        ff         381             22   NaN
15        ff         867            381  flag # 'flag' row &
16        ff        from            867   NaN # 'from' row are adjacent for 'ff'

如果有人能提供帮助，真的很感激。

【问题讨论】：

标签： python pandas dataframe loops group-by

【解决方案1】：

让我们groupby custom_id 和 filter 组上的数据框使用自定义 lambda 函数 f 根据指定条件返回布尔值：

f = lambda g: (g['sending_num'].eq('from') & g['note'].shift().eq('flag')).any()
sub_df = df.groupby('custom_id').filter(f)

或者您可以先根据指定条件创建一个布尔掩码，然后使用此掩码获取满足规则的custom_id：

m = df.groupby('custom_id')['note'].shift().eq('flag') & df['sending_num'].eq('from')
sub_df = df[df['custom_id'].isin(df.loc[m, 'custom_id'].unique())].copy()

print(sub_df)

   custom_id sending_num  receiving_num  note
0         aa          11            900   NaN
1         aa         252             11  flag
2         aa         198            252   NaN
3         aa         266            198   NaN
4         aa        5317            266  flag
5         aa        from           5317   NaN
11        dd          22            432   NaN
12        dd          55             22  flag
13        dd        from             55   NaN
14        dd         376             65   NaN
15        dd          98            376   NaN
16        ff         901             42   NaN
17        ff         126            901   NaN
18        ff          22            126   NaN
19        ff         381             22   NaN
20        ff         867            381  flag
21        ff        from            867   NaN

【讨论】：