Python Pandas - Calculate mean using criteria from two columns

Question

I am trying to calculate the mean of a Pandas data frame column using selection criteria from two other columns. In the code below, there are a "Trace" and "Sample" columns that are to be used for selection criteria, while the numbers in the "Value" column are to be used in the calculation. I want to group by the "Trace" number and only take the average of "Sample" numbers 3, 4, and 5. Then, I would like to create a new column in the original dataframe "df" and place the calculated mean values in all of the rows corresponding to the correct "Trace" number.

d = {"Trace": [1,1,1,1,1,2,2,2,2,2], "Sample": [1,2,3,4,5,1,2,3,4,5], "Value": [2,3,5,6,1,8,9,5,4,3]}

Any ideas?

Thanks!

Answer 1

You can try this, filter your dataframe first, then groupy with mean, and join back to original dataframe on 'Trace' (which is the common column name between the dataframes are reset_index on results of groupby):

df[df['Sample'].isin([3,4,5])].groupby('Trace')['Value'].mean()\
                              .rename('Avg Value').reset_index().merge(df)

Output:

   Trace  Avg Value  Sample  Value
0      1          4       1      2
1      1          4       2      3
2      1          4       3      5
3      1          4       4      6
4      1          4       5      1
5      2          4       1      8
6      2          4       2      9
7      2          4       3      5
8      2          4       4      4
9      2          4       5      3

OR

df.groupby('Trace')\
  .apply(lambda x: x.loc[x['Sample'].isin([3,4,5]),'Value'].mean())\
  .rename('Avg Value').reset_index().merge(df)

Output:

   Trace  Avg Value  Sample  Value
0      1        4.0       1      2
1      1        4.0       2      3
2      1        4.0       3      5
3      1        4.0       4      6
4      1        4.0       5      1
5      2        4.0       1      8
6      2        4.0       2      9
7      2        4.0       3      5
8      2        4.0       4      4
9      2        4.0       5      3