filtering rows after a groupby and apply a function

user_13070745 0
Viewed 64

I'm using python and pandas to work on some data. My data looks like the following:

df = pd.DataFrame({'A' : ['foo', 'bar', 'foo', 'bar',
                         'foo', 'bar'],
                   'B' : [1, 2, 3, 4, 5, 6],
                   'C' : [True, False, True, True, False, True]})
print(df)

     A  B      C
0  foo  1   True
1  bar  2  False
2  foo  3   True
3  bar  4   True
4  foo  5  False
5  bar  6   True

What I would like to do:

  1. Groupby by "A"
  2. Select the value B by groups where C == True
  3. Calculate a mean value on this selection
  4. Create a new column "D" to store theses values

So the result would be:

    A   B   C       D
0   foo 1   True    2
1   bar 2   False   5
2   foo 3   True    2
3   bar 4   True    5
4   foo 5   False   2
5   bar 6   True    5

I have tried some mixes of groupby, filter and transform, but I cannot succed to make it work. I imagine the solution close to the followings

df.groupby(["A"])[df.loc[df["C"] == True, "B"]].transform("mean")

or

df.groupby(["A"]).filter(lambda x: x["D"] == True)["B"].transform("mean")

But none of these syntax work.

Thanks for helping me and people in general,

pandas-groupby pandas python
1 Answers

Use Series.map with means of filtered rows, ==True should be omitted:

df['D'] = df['A'].map(df.loc[df.C, 'B'].groupby(df["A"]).mean())
print(df)

     A  B      C  D
0  foo  1   True  2
1  bar  2  False  5
2  foo  3   True  2
3  bar  4   True  5
4  foo  5  False  2
5  bar  6   True  5
user_29010020
Related