Python的大熊貓：基於對另外兩列

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 

我可以選擇的values總和每當variableA使用variableA當前行相匹配：

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

，但選擇的values每當variableB匹配的variableA當前行逃避我的總和。我嘗試了.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)) 

---

時間測試

Answer 2

嗯，你總是可以使用.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