Getting corresponding values in a groupby

user_12957999 0
Viewed 71

I have a dataset similar to this

Serial    A        B       
1         12               
1         31
1
1         12
1         31       203
1         10    
1         2
2         32       100     
2         32       242
2         3
3         2                
3         23       100     
3
3         23

I group the dataframe based on Serial and find the maximum value of each A column by df['A_MAX'] = df.groupby('Serial')['A'].transform('max').values and retain the first value by df['A_MAX'] = df['A_MAX'].mask(df['Serial'].duplicated(), '')

Serial    A        B       A_MAX    B_corresponding
1         12               31       203
1         31
1
1         12
1         31       203
1         10    
1         2
2         32       100     32       100
2         32       242
2         3
3         2                23       100
3         23       100     
3
3         23

Now for the B_corresponding column, I would like to get the corresponding B values of the A_MAX. I thought of locating the A_MAX values in A but there are similar max A values per group. Additional condition, for example in Serial 2 I would also prefer to get the smallest B values between the 32

pandas python
2 Answers

Idea is use DataFrame.sort_values for maximal values per groups, then remove missing values by DataFrame.dropna and get first rows by Serial by DataFrame.drop_duplicates. Create Series by DataFrame.set_index and last use Series.map:

df['A_MAX'] = df.groupby('Serial')['A'].transform('max')
df['A_MAX'] = df['A_MAX'].mask(df['Serial'].duplicated())

s = (df.sort_values(['Serial','A'], ascending=[True, False])
       .dropna(subset=['B'])
       .drop_duplicates('Serial')
       .set_index('Serial')['B'])
df['B_corresponding'] = df['Serial'].map(s).mask(df['Serial'].duplicated())
print (df)
    Serial     A      B  A_MAX  B_corresponding
0        1  12.0    NaN   31.0            203.0
1        1  31.0    NaN    NaN              NaN
2        1   NaN    NaN    NaN              NaN
3        1  12.0    NaN    NaN              NaN
4        1  31.0  203.0    NaN              NaN
5        1  10.0    NaN    NaN              NaN
6        1   2.0    NaN    NaN              NaN
7        2  32.0  100.0   32.0            100.0
8        2  32.0  242.0    NaN              NaN
9        2   3.0    NaN    NaN              NaN
10       3   2.0    NaN   23.0            100.0
11       3  23.0  100.0    NaN              NaN
12       3   NaN    NaN    NaN              NaN
13       3  23.0    NaN    NaN              NaN

Converting missing values to empty strings is possible, but get mixed values - numeric and strings, so next processing should be problematic:

df['A_MAX'] = df.groupby('Serial')['A'].transform('max')
df['A_MAX'] = df['A_MAX'].mask(df['Serial'].duplicated(), '')

s = (df.sort_values(['Serial','A'], ascending=[True, False])
       .dropna(subset=['B'])
       .drop_duplicates('Serial')
       .set_index('Serial')['B'])

df['B_corresponding'] = df['Serial'].map(s).mask(df['Serial'].duplicated(), '')
print (df)
    Serial     A      B A_MAX B_corresponding
0        1  12.0    NaN    31             203
1        1  31.0    NaN                      
2        1   NaN    NaN                      
3        1  12.0    NaN                      
4        1  31.0  203.0                      
5        1  10.0    NaN                      
6        1   2.0    NaN                      
7        2  32.0  100.0    32             100
8        2  32.0  242.0                      
9        2   3.0    NaN                      
10       3   2.0    NaN    23             100
11       3  23.0  100.0                      
12       3   NaN    NaN                      
13       3  23.0    NaN
user_29010020

You could also use dictionaries to achieve the same if you are not so inclined to only use pandas.

a_to_b_mapping = df.groupby('A')['B'].min().to_dict()
series_to_a_mapping = df.groupby('Series')['A'].max().to_dict()
agg_df = {}
for series, a in series_to_a_mapping.items():
    agg_df.append((series, a, a_to_b_mapping.get(a, None)))
  
agg_df = pd.DataFrame(agg_df, columns=['Series', 'A_max', 'B_corresponding'])
agg_df.head()

    Series  A_max   B_corresponding
0   1   31.0    203.0
1   2   32.0    100.0
2   3   23.0    100.0

If you want, you could join this to original dataframe and mask duplicates.

dft = df.join(final_df.set_index('Serial'), on='Serial', how='left')
dft['A_max'] = dft['A_max'].mask(dft['A_max'].duplicated(), '')
dft['B_corresponding'] = dft['B_corresponding'].mask(dft['B_corresponding'].duplicated(), '')
dft
user_96479510
Related