Look if element is present for each ID after groupby

user_13688384 0
Viewed 43

Here is my dataframe:

       ID  Color 
0       1   red    
1       1   blue
2       1   yellow
3       2   blue
4       2   purple
5       3   yellow
6       3   green

I want to create a third column that tells me whether there is a color red or yellow for each ID. If there is a red then the third column will be 1, if there is a purple then the third column will be 0, if there are neither it stays blank. There will be no cases that red and purple will be in the same ID.

So, my resultant df should look like this:

       ID  Color    col3
0       1   red     1
1       1   blue    1
2       1   yellow  1
3       2   blue    0
4       2   purple  0
5       3   yellow  
6       3   green

I'm thinking about doing a groupby on my ID column, then checking if there is a red or purple present in the color column. I'm not sure how to implement this.

Please help

pandas-groupby pandas python
2 Answers

Here's one way using np.select:

conds = [df.Color.eq('red'), df.Color.eq('purple')]
df['col3'] = np.select(conds, [1,0], '')
df['col3'] = df.groupby('ID').col3.transform('max')

Or we could instead set a nan as default, and transform with first:

df['col3'] = np.select(conds, [1,0], np.nan)
df['col3'] = df.groupby('ID').col3.transform('first').fillna('')

print(df)

   ID   Color col3
0   1     red    1
1   1    blue    1
2   1  yellow    1
3   2    blue    0
4   2  purple    0
5   3  yellow     
6   3   green

Note that the former approach takes advantage of the following:

max('', '0')
# '0'
max('', '1')
# '1'
user_96986840

Here's an alternative if you don't want to use numpy for whatever reason:

df['col3'] = df.set_index('ID')['Color'].apply({'red': 1, 'purple': 0}.get).groupby(level=0).transform('max').fillna('').reset_index(drop=True)

   ID   Color col3
0   1     red    1
1   1    blue    1
2   1  yellow    1
3   2    blue    0
4   2  purple    0
5   3  yellow     
6   3   green
user_91833440
Related