Iterate over a groupby dataframe to operate in each row

Viewed 49

I have a DataFrame like this:

    subject  trial  attended
0         1      1         1
1         1      3         0
2         1      4         1
3         1      7         0
4         1      8         1
5         2      1         1
6         2      2         1
7         2      6         1
8         2      8         0
9         2      9         1
10        2     11         1
11        2     12         1
12        2     13         1
13        2     14         1
14        2     15         1
  1. I would like to GroupBy subject.
  2. Then iterate in each row of the GroupBy dataframe.
  3. If for a row 'attended' == 1, then to increase a variable sum_reactive by 1.
  4. If the sum_reactive variable reaches == 4, then to add in a dictionary the 'subject' and 'trial' in which the variable sum_reactive reached a value of 4.

I as trying to define a function for this, but it doesn't work:

def count_attended():
    sum_reactive = 0
    dict_attended = {}
    for i, g in reactive.groupby(['subject']):
        for row in g:
            if g['attended'][row] == 1:
                sum_reactive += 1
                if sum_reactive == 4:
                   dict_attended.update({g['subject'] : g['trial'][row]})
                   return dict_attended

    return dict_attended

I think that I don't have clear how to iterate inside each GroupBy dataframe. I'm quite new using pandas.

1 Answers

IIUC try,

df = df.query('attended == 1')
df.loc[df.groupby('subject')['attended'].cumsum() == 4, ['subject', 'trial']].to_dict(orient='record')

Output:

[{'subject': 2, 'trial': 9}]

Using groupby with cumsum will do the counting attended, then check to see when this value equals to 4 to create a boolean series. You can use this boolean series to do boolean indexing to filter your dataframe to certain rows. Lastly, with lock and column filtering select subject and trial.

Related