Python Pandas：基于其他两列的值查找列总和

Question

在遍历variableA列时，我想生成一个新的列，该列是values之和，只要任一variableA或variableB中的行等于variableA当前行的值。示例数据：

    values    variableA  variableB
  0  134       1             3
  1  12        2             6
  2  43        1             2
  3  54        3             1
  4  16        2             7

当variableA与当前行的variableA匹配时，我可以使用以下代码选择values的总和：

df.groupby('variableA')['values'].transform('sum')

但是每当variableB匹配当前行的variableA时，选择values的总和却让我感到困惑。我尝试过使用.loc，但似乎与.groupby不太兼容。期望的输出应该如下所示：

    values    variableA  variableB  result
  0  134       1             3      231
  1  12        2             6      71
  2  43        1             2      231
  3  54        3             1      188
  4  16        2             7      71

谢谢！

- Flow Nuwen

Answer 1

嗯，你总是可以使用.apply，但要注意：它可能会很慢：

>>> df
   values  variableA  variableB
0     134          1          3
1      12          2          6
2      43          1          2
3      54          3          1
4      16          2          7
>>> df.apply(lambda S: df.loc[(df.variableA == S.variableA) | (df.variableB == S.variableA), 'values'].sum(), axis=1)
0    231
1     71
2    231
3    188
4     71
dtype: int64

当然，您需要分配它...

>>> df['result'] = df.apply(lambda S: df.loc[(df.variableA == S.variableA) | (df.variableB == S.variableA), 'values'].sum(), axis=1)
>>> df
   values  variableA  variableB  result
0     134          1          3     231
1      12          2          6      71
2      43          1          2     231
3      54          3          1     188
4      16          2          7      71

- juanpa.arrivillaga

Answer 2

使用numpy广播的向量化方法

vars = df[['variableA', 'variableB']].values
matches = (vars[:, None] == vars[:, [0]]).any(-1)

df.assign(result=df['values'].values @ matches)  # @ operator with python 3
# use this for use python 2
# df.assign(result=df['values'].values.dot(matches))

---

时间测试

- piRSquared