I have 2 dataframes like this...
np.random.seed(0)
a = pd.DataFrame(np.random.randn(20,3))
b = pd.DataFrame(np.random.randint(1,5,size=(20,3)))
I'd like to find the average of values in a for the 4 groups in b.
This...
a[b==1].sum().sum() / a[b==1].count().sum()
...works for doing one group at a time, but I was wondering if anyone could think of a cleaner method.
My expected result is
1 -0.088715
2 -0.340043
3 -0.045596
4 0.582136
dtype: float64
Thanks.
If you want a fast numpy solution, use np.unique and np.bincount:
c, d = (a_.to_numpy().ravel() for a_ in [a, b])
u, i, cnt = np.unique(d, return_inverse=True, return_counts=True)
np.bincount(i, c) / cnt
# array([-0.0887145 , -0.34004319, -0.04559595, 0.58213553])
To construct a Series, use
pd.Series(np.bincount(i, c) / cnt, index=u)
1 -0.088715
2 -0.340043
3 -0.045596
4 0.582136
dtype: float64
For comparison, stack returns,
a.stack().groupby(b.stack()).mean()
1 -0.088715
2 -0.340043
3 -0.045596
4 0.582136
dtype: float64
%timeit a.stack().groupby(b.stack()).mean()
%%timeit
c, d = (a_.to_numpy().ravel() for a_ in [a, b])
u, i, cnt = np.unique(d, return_inverse=True, return_counts=True)
np.bincount(i, c) / cnt
5.16 ms ± 305 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
113 µs ± 1.92 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)