pandas: select sub-datasets by complex condition in each group

user_12136878 0
Viewed 99

I need to select a subset from the given dataframe. Here is the 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]
                   })

And the df is like this:

   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

I wish to select a subset according to the following rules: for each group (custom id), if: there is a 'from' appearing in the row, AND in its 'note' column, the row above has a 'flag' value. For example, for group 'aa', there is a 'from' in its 'sending_num' column, in the meanwhile, in its row above (row 4) there is a 'flag' in the 'note' column in the same group, so 'aa' is a target one; Similar like group 'dd' and 'ff' because in their 'sending_num' column there is 'from', and in the row above in the 'note' column, there is 'flag', so these two are selected, but not the other groups. I tried to write a loop and iloc to do this, but very slow. Eventually, I wish to have a subset like this according to the rules:

   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'

Really appreciate it if anyone can help.

pandas dataframe group-by python loops
1 Answers

Let's groupby the dataframe on custom_id and filter the groups with a custom lambda function f which returns a boolean value based on the specified condition:

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

Alternatively you can first create a boolean mask based on the specified condition, then use this mask to get the custom_id that satisfy the rule:

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
user_128331660
Related