What's the most efficient way to iterate by rows for each group of rows?

Viewed 79

I'm wondering how to efficiently loop through rows by groups. So like the following sample dataset shows, it includes 3 different students with their pass records in 3 months.

import pandas as pd
import numpy as np
df = pd.DataFrame({'student':'A A A B B B C C C'.split(),
                  'month':[1, 2, 3, 1, 2, 3, 1, 2, 3],
                  'pass':[0, 1, 0, 0, 0, 0, 1, 0, 0]})
print(df)
 student  month  pass
0       A      1     0
1       A      2     1
2       A      3     0
3       B      1     0
4       B      2     0
5       B      3     0
6       C      1     1
7       C      2     0
8       C      3     0

I'd like to have a new column "pass_patch", which should be equal to "pass" at first. But when a student has "pass" as 1 then all of his "pass_patch" in the following months should be 1, like the following:

df = pd.DataFrame({'student':'A A A B B B C C C'.split(),
                   'month':[1, 2, 3, 1, 2, 3, 1, 2, 3],
                   'pass':[0, 1, 0, 0, 0, 0, 1, 0, 0],
                   'pass_patch':[0, 1, 1, 0, 0, 0, 1, 1, 1]})
print(df)
  student  month  pass  pass_patch
0       A      1     0           0
1       A      2     1           1
2       A      3     0           1
3       B      1     0           0
4       B      2     0           0
5       B      3     0           0
6       C      1     1           1
7       C      2     0           1
8       C      3     0           1

I did some searches and found iterrows might work, but was concerned it would be too slow to run for the whole dataset (around million of records). Would there be more efficient ways to realize that?

Any suggestions would be greatly appreciated.

4 Answers

Try with cummax

df['new'] = df.groupby('student')['pass'].cummax()
df
Out[78]: 
  student  month  pass  new
0       A      1     0    0
1       A      2     1    1
2       A      3     0    1
3       B      1     0    0
4       B      2     0    0
5       B      3     0    0
6       C      1     1    1
7       C      2     0    1
8       C      3     0    1

What's the most efficient way to iterate by rows for each group of rows

DON'T ITERATE MANUALLY

Manualy iteration should always be your last option to try, most often, there's always some better way to perfrom a required operation rather than doing the iteration.

You can groupby student, then call cumsum which will just sum the values iteratively, convert it to boolean then convert it back to int:

df['pass_patch'] = df.groupby('student')['pass'].cumsum().astype(bool).astype(int)

OUTPUT:

  student  month  pass  pass_patch
0       A      1     0           0
1       A      2     1           1
2       A      3     0           1
3       B      1     0           0
4       B      2     0           0
5       B      3     0           0
6       C      1     1           1
7       C      2     0           1
8       C      3     0           1

PS: In above solution, you can avoid .astype(bool).astype(int) part if there is not more than one 1s for pass for each group. You may also need to sort the dataframe on months for each student if they are not sorted, I have not added that part since the sample data you have provided is already in that order.

You can replace 0 by pd.NA and then use ffill method, and then replace the null values back to 0:

df['pass_patch'] = df['pass'].replace(0, pd.NA)
df['pass_patch'] = df.groupby('student')['pass_patch']\
                    .transform(lambda x: x.ffill())\
                    .fillna(0)\
                    .astype(int)

Output:

  student  month  pass  pass_patch
0       A      1     0           0
1       A      2     1           1
2       A      3     0           1
3       B      1     0           0
4       B      2     0           0
5       B      3     0           0
6       C      1     1           1
7       C      2     0           1
8       C      3     0           1
import pandas as pd
import numpy as np

df = pd.DataFrame({'student':'A A A B B B C C C'.split(),
                  'month':[1, 2, 3, 1, 2, 3, 1, 2, 3],
                  'pass':[0, 1, 0, 0, 0, 0, 1, 0, 0]})

First we search the month of first pass for students that actually passed at least once.

grp = df[df["pass"].eq(1)]\
    .sort_values(["student", "month"])\
    .groupby("student").head(1)

where grp looks like

  student  month  pass
1       A      2     1
6       C      1     1

Then we merge the dataframes

df = pd.merge(df, 
              grp,
              on=["student"],
              how="left",
              suffixes=(None, '_y'))

and df looks like

  student  month  pass  month_y  pass_y
0       A      1     0      2.0     1.0
1       A      2     1      2.0     1.0
2       A      3     0      2.0     1.0
3       B      1     0      NaN     NaN
4       B      2     0      NaN     NaN
5       B      3     0      NaN     NaN
6       C      1     1      1.0     1.0
7       C      2     0      1.0     1.0
8       C      3     0      1.0     1.0

Finally we set 1 to all months greater or equals to month_y and 0 otherwise.

df["pass_patch"] = np.where(
    df["month"].ge(df["month_y"]),
    1,
    0)

and we drop the columns we don't need anymore

df = df.drop(columns=["month_y", "pass_y"])

Which returns

  student  month  pass  pass_patch
0       A      1     0           0
1       A      2     1           1
2       A      3     0           1
3       B      1     0           0
4       B      2     0           0
5       B      3     0           0
6       C      1     1           1
7       C      2     0           1
8       C      3     0           1
Related