Pandas：条件分组特定计算

Question

假设我有一个表格，里面有一个关键字段（比如客户ID）和两个数字列C1和C2。我想根据这个关键字段（客户）来分组这些行，然后对这些列进行一些汇总，比如求和和平均值。计算完分组汇总后，我想把结果再加回到每个客户的行中，作为一些客户的特征。

我可以这样做：
df['F1'] = df.groupby(['Key'])['C1'].transform(np.sum)
如果我只想对一列进行汇总，并且能够把结果加回到数据框中，这样是可以的。

那么我能不能设置条件呢？比如我只想对C2列等于某个数字X的行进行C1列的求和，并且还能把结果加回到数据框中？

我该如何对一组行进行汇总，比如：
np.sum(C1 + C2)？

实现这个功能最简单、最优雅的方法是什么？最有效率的方法又是什么？这些汇总能不能在一个步骤中完成？

谢谢大家的帮助。

Answer 1

我找到了一种不同的方法，它使用apply()而不是transform()，不过你需要把结果表和输入的数据表连接起来，我还没弄明白怎么做。希望能有人帮我完成这个连接表的部分，或者给我一些更好的建议。

df = pd.DataFrame({'Key': ['a','a','b','b','c','c'],
                            'C1': [1,2,3,4,5,6],
                            'C2': [7,8,9,10,11,12]})

# Group g will be given as a DataFrame
def group_feature_extractor(g):
    feature_1 = (g['C1'] + g['C2']).sum()

    even_C1_filter = g['C1'] % 2 == 0
    feature_2 = g[even_C1_filter]['C2'].sum()

    return pd.Series([feature_1, feature_2], index = ['F1', 'F2'])

# Group once
group = df.groupby(['Key'])

# Extract features from each group
group_features = group.apply(group_feature_extractor)

#
# Join with the input data frame ...
#

Answer 2

这里是一些虚拟数据的设置。

In [81]: df = pd.DataFrame({'Key': ['a','a','b','b','c','c'], 
                            'C1': [1,2,3,4,5,6],  
                            'C2': [7,8,9,10,11,12]})
In [82]: df['F1'] = df.groupby('Key')['C1'].transform(np.sum)

In [83]: df
Out[83]: 
   C1  C2 Key  F1
0   1   7   a   3
1   2   8   a   3
2   3   9   b   7
3   4  10   b   7
4   5  11   c  11
5   6  12   c  11

如果你想要进行条件分组，可以在传递给 .groupby 的时候先过滤一下数据框。例如，如果你想要计算 'C1' 的组总和，但只在 'C2' 小于 8 或大于 9 的情况下。

In [87]: cond = (df['C2'] < 8) | (df['C2'] > 9)

In [88]: df['F2'] = df[cond].groupby('Key')['C1'].transform(np.sum)

In [89]: df
Out[89]: 
   C1  C2 Key  F1  F2
0   1   7   a   3   1
1   2   8   a   3 NaN
2   3   9   b   7 NaN
3   4  10   b   7   4
4   5  11   c  11  11
5   6  12   c  11  11

这样做是可行的，因为 transform 操作会保留索引，所以它仍然会和原始数据框正确对齐。

如果你想要对两个列的组总和进行求和，可能这样做会比较简单？也许有人有更聪明的方法。

In [93]: gb = df.groupby('Key')

In [94]: df['C1+C2'] = gb['C1'].transform(np.sum) + gb['C2'].transform(np.sum)

编辑：还有一种方法可以获取多个列的组总和。语法上并没有更简洁，但对于很多列来说可能会更方便。

df['C1_C2'] = gb[['C1','C2']].apply(lambda x: pd.DataFrame(x.sum().sum(), index=x.index, columns=['']))