根据多个组合条件创建新列

Question

I want a new column in this df with the following condition. The column education is a categorical value that goes from 1 to 5 (1 is the lower level of education and 5 is the higher level of education). I want to create a function with the following logic (so as to create a new column in the df)

First, for any id check if there is at least a education level graduated, then the new column must have the higher level of education graduated.

Second, if there is no graduated education level for some particular id (must have all educaction level in "In course"). So, must check the maximium level of education and substract one.

df
id  education stage
1   2         Graduated
1   3         Graduated
1   4         In course
2   3         In course
3   2         Graduated
3   3         In course
4   2         In course

expected output:

id  education stage       new_column
1   2         Graduated   3
1   3         Graduated   3
1   4         In course   3
2   3         In course   2
3   2         Graduated   2
3   3         In course   2
4   2         In course   1

Answer 1

You can do it like this:

import pandas as pd
df = pd.DataFrame({'id': [1, 1, 1, 2, 3, 3, 4], 'education': [2, 3, 4, 3, 2, 3, 2],
                   'stage': ['Graduated', 'Graduated', 'In course', 'In course', 'Graduated', 'In course', 'In course']})


max_gr = df[df.stage == 'Graduated'].groupby('id').education.max()
max_ic = df[df.stage == 'In course'].groupby('id').education.max()

# set all cells to the value from max_ed
df['new_col'] = df.id.map(max_gr)
# set cells that have not been filled to the value from max_ic - 1
df.loc[df.new_col.isna(), ['new_col']] = df.id.map(max_ic - 1)

series.map(other_series) returns a new series where the values from series have been replaced by the values from other_series .

Answer 2

This is one way.

df['new'] = df.loc[df['stage'] == 'Graduated']\
              .groupby('id')['education']\
              .transform(max).astype(int)

df['new'] = df['new'].fillna(df.loc[df['stage'] == 'InCourse']\
                               .groupby('id')['education']\
                               .transform(max).sub(1)).astype(int)

Result

   id  education      stage  new
0   1          2  Graduated    3
1   1          3  Graduated    3
2   1          4   InCourse    3
3   2          3   InCourse    2
4   3          2  Graduated    2
5   3          3   InCourse    2
6   4          2   InCourse    1

Explanation

• First, map to "Graduated" dataset grouped by id on max education.
• Second, map to "InCourse" dataset grouped by id on max education minus 1.

Answer 3

Alternative solution based on Markus Löffler.

max_ic = df[df.stage.eq('In course')].groupby('id').education.max() - 1
max_gr = df[df.stage.eq('Graduated')].groupby('id').education.max()

# Update with max_gr
max_ic.update(max_gr)

df['new_col'] = df.id.map(max_ic)