Home > OS >  pandas groupby apply with condition on the first occurrence of a column value
pandas groupby apply with condition on the first occurrence of a column value

Time:11-22

I have a data frame shown below with pid and event_date being the indices after applying groupby. I want to apply groupby again this time only to pid, and applies to two conditions:

  1. A person (pid=person) has two or more True labels;
  2. The first True instance of this person occurred when he/she was under 45 years old;

If the two above conditions satisfy then assign this person/pid to True in the groupby-ed dataframe.

                           age      label
  pid       event_date      
00000001    2000-08-28  76.334247   False
            2000-10-17  76.471233   False
            2000-10-31  76.509589   True
            2000-11-02  76.512329   True
... ... ... ...
00000005    2014-08-15  42.769863   False
            2015-04-04  43.476712   False
            2015-11-06  44.057534   True
            2017-03-06  45.386301   True

I have come only so far to implement the first condition:

df = (df.groupby(['pid']).apply(lambda x: sum(x['label'])>1).to_frame('label'))

The second one is tricky for me. How do I condition on the first occurrence of some column value? Any advice is very much welcomed! Many thanks!

UPDATE with an example dataframe:

a = pd.DataFrame(columns=['pid', 'event_date', 'age', 'label'])
a['pid'] = [1, 1, 1, 1, 5, 5, 5, 5]
a['event_date'] = ['2000-08-28', '2000-08-28', '2000-08-28', '2000-08-28',\
                  '2000-08-28', '2000-08-28', '2000-08-28', '2000-08-28']
a['event_date'] = pd.to_datetime(a.event_date)
a['age'] = [76.334247, 76.471233, 76.509589, 76.512329, 42.769863, 43.476712, 44.057534, 45.386301]
a['label'] = [False, False, True, True, False, False, True, True]

a = (a.groupby(['pid', 'event_date', 'age']).apply(lambda x: x['label'].any()).to_frame('label'))
a.reset_index(level=['age'], inplace=True)

Now if I apply (a.groupby(['pid']).apply(lambda x: sum(x['label'])>1).to_frame('label')) I would get

    label
pid 
1   True
5   True

Which only satisfies the first condition (well because I skipped the second one). Adding the second condition should only label pid=5 True since only this person/pid was under 45 when the first label=True occurred.

CodePudding user response:

After half a (fun) hour, I came up with this:

condition = a.reset_index().groupby('pid')['label'].sum().ge(2) & a.reset_index().groupby('pid').apply(lambda x: x['age'][x['label'].idxmax()] < 45)

Output:

>>> condition
pid
1    False
5     True
dtype: bool

It could be shorten a little bit (removing the two .reset_index() calls) if the index was normal, not a MultiIndex of pid event_date. If you can't avoid that from the start and you don't mind changing a:

a = a.reset_index()
condition = a.groupby('pid')['label'].sum().ge(2) & a.groupby('pid').apply(lambda x: x['age'][x['label'].idxmax()] < 45)

Expanded:

condition = (
    a.groupby('pid') # Group by pid
    ['label']        # Get the label column for each group
    .sum()           # Compute the sum of the True values
    .ge(2)           # Are there two or more?
    
    & # Boolean mask. The previous and the next bits of code are the two conditions, and they return a series, where the index is each unique pid, and the value is whether the condition is met for all the rows in that pid
    
    a.groupby('pid')                # Group by pid
    .apply(                         # Call a function for each group, passing the group (a dataframe) to the function as its first parameter
        lambda x:                   # Function start
            x['age'][               # Get item from the age column at the specified index
                x['label'].idxmax() # Get the index of the highest value of the label column (since they're only boolean values, the highest will be the first True value)
            ] < 45                  # Check if it's less than 45
    )
)
  • Related