根据组计算dataframe列

时间:2017-05-02 16:14:42

标签: python pandas math

我有一个数据框,我想为其计算一个新列。该列将包含 x + b y; x和y是列,而a和b是标量。问题在于标量系数对于第三列“组”中定义的不同组而言是不同的。我尝试通过循环遍历组来做到这一点,但我得到NULL作为我的输出。 “系数”是列表[a,b]

的字典
for group in df['groups'].unique():
    coef_list = Coefficients[group]
    temp = df[df['groups']==group ]
    df['calculation'] = coef_list[0] * temp['x'] + coef_list[1] * temp['y'] 

我认为仅通过匹配索引来添加值就可以了。我是否需要使用assign或其他内容。

1 个答案:

答案 0 :(得分:2)

考虑数据框df和系数字典coef

coef = dict(G1=[1, 2], G2=[3, 4])
df = pd.DataFrame(dict(
        groups=['G1'] * 4 + ['G2'] * 4,
        x=range(1, 9), y=range(1, 9)[::-1]
    ))

<强> 解决方案

我们可以使用map生成分组值并使用pd.DataFrame构造函数

myvars = ['x', 'y']
gvals = pd.DataFrame(df.groups.map(coef).values.tolist(), df.index, myvars)
gvals

   x  y
0  1  2
1  1  2
2  1  2
3  1  2
4  3  4
5  3  4
6  3  4
7  3  4

我们可以在乘法和求和后assign新列。

df.assign(calculation=gvals.mul(df[myvars]).sum(1))

  groups  x  y  calculation
0     G1  1  8           17
1     G1  2  7           16
2     G1  3  6           15
3     G1  4  5           14
4     G2  5  4           31
5     G2  6  3           30
6     G2  7  2           29
7     G2  8  1           28

通过访问numpy数组

可以更快地做到这一点
df.assign(calculation=(gvals.values * df[myvars].values).sum(1))