熊猫按 2 列分组，应用函数，选择最大值并返回索引值

Question

这是我正在尝试执行的操作：

ID    SUB_ID    AMOUNT
1   101     1        50
2   101     1        -10
3   101     1        -20
4   101     2        30
5   101     2        20
6   102     3        10
7   102     3        -10
8   102     4        10
9   102     4        10

我们要按ID和SUB_ID分组，然后取AMOUNT的绝对值之和。然后在ID组中对这个汇总列进行排序，并返回最大值的SUB_ID值。

我们可以通过以下方式获得求和：

df1 = (df
.groupby(['ID','SUB_ID'])
.apply(lambda x: np.sum(np.absolute(x['AMOUNT']))))
)

这将返回一个具有多索引的序列

ID    SUB_ID    
101     1        80
2        50
102     3        20
4        20

从这里我想返回 [1,3]([1,4] 也被接受，因为 102 组中的两个值是相同的，但我们只想为每个组返回一个值！

显然，我们可以循环并选择最大值，但我正在努力找出最有效的方法。此操作将应用于数百万行。

Answer 1

这是一种方式。由于您的数据集很大，我强烈建议您避免使用lambda函数，因为它们不是以矢量化方式应用的。

res = df.assign(AMOUNT=df['AMOUNT'].abs())
.groupby(['ID', 'SUB_ID'], as_index=False).sum()
.sort_values('AMOUNT', ascending=False)
.groupby('ID').head(1)

---

例

df = pd.DataFrame([[101, 1, 50], [101, 1, -10], [101, 1, -20], [101, 2, 30],
[101, 2, 20], [102, 3, 10], [102, 3, -10], [102, 4, 10], [102, 4, 10]],
columns=['ID', 'SUB_ID', 'AMOUNT'])
res = df.assign(AMOUNT=df['AMOUNT'].abs())
.groupby(['ID', 'SUB_ID'], as_index=False).sum()
.sort_values('AMOUNT', ascending=False)
.groupby('ID').head(1)
print(res)
ID  SUB_ID  AMOUNT
0  101       1      80
2  102       3      20

Answer 2

我认为您可以使用nlargest：

df1.groupby('ID').nlargest(1).index.get_level_values(level='SUB_ID').tolist()
# [1, 3]