python Pandas : Find Sum of Column Based on Value of Two other Columns

Question

在遍历 variableA 列时，我想生成一个新列，它是 values 的总和，每当 either 中的一行 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

谢谢!

Answer 1

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

---

时间测试