Python Pandas：根据其他两列的值查找列的总和

Question

在遍历variableA 列时，我想生成一个新列，只要variableA 或variableB 或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

谢谢！

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

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))

---

时间测试