Pandas 聚合 - 如何保留所有列

Question

下面是一个示例数据表：

rand = np.random.RandomState(1)
df = pd.DataFrame({'A': ['group1', 'group2', 'group3'] * 2,
                'B': rand.rand(6),
                'C': rand.rand(6),
                'D': rand.rand(6)})

打印数据表 df

        A         B         C         D
0  group1  0.417022  0.186260  0.204452
1  group2  0.720324  0.345561  0.878117
2  group3  0.000114  0.396767  0.027388
3  group1  0.302333  0.538817  0.670468
4  group2  0.146756  0.419195  0.417305
5  group3  0.092339  0.685220  0.558690

按 A 列进行分组

group = df.groupby('A')

使用 agg 函数来返回每组的最大值

max1 = group['B'].agg({'max' : np.max})
print max1

             max
A               
group1  0.417022
group2  0.720324
group3  0.092339

但是我想保留（或者找回）其他列 C 和 D 中的相关数据。这些数据是包含最大值的那一行的剩余数据。

所以，返回的结果应该是：

     A         B         C         D
group1  0.417022  0.186260  0.204452
group2  0.720324  0.345561  0.878117
group3  0.092339  0.685220  0.558690

有没有人能告诉我怎么做？任何帮助都非常感谢。

Answer 1

我的回答和FooBar的差不多，不过我用idmax()这个函数把它简化成了一行代码。

df.loc[df.groupby('A')['B'].idxmax()]

结果是一样的：

In [51]: df
Out[51]: 
        A         B         C         D
0  group1  0.417022  0.186260  0.204452
1  group2  0.720324  0.345561  0.878117
2  group3  0.000114  0.396767  0.027388
3  group1  0.302333  0.538817  0.670468
4  group2  0.146756  0.419195  0.417305
5  group3  0.092339  0.685220  0.558690

In [76]: df.loc[df.groupby('A')['B'].idxmax()]
Out[76]: 
        A         B         C         D
0  group1  0.417022  0.186260  0.204452
1  group2  0.720324  0.345561  0.878117
5  group3  0.092339  0.685220  0.558690

Answer 2

这个过程分为两个步骤：第一步是找到索引，第二步是查找所有的行。

idx = df.groupby('A').apply(lambda x: x['B'].argmax())
idx

Out[362]: 
A
group1    0
group2    1
group3    5

df.loc[idx]

Out[364]: 
        A         B         C         D
0  group1  0.417022  0.186260  0.204452
1  group2  0.720324  0.345561  0.878117
5  group3  0.092339  0.685220  0.558690