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

Question

Here is the operation I am trying to do:

    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

We want to group by ID and SUB_ID , and then take the sum of the absolute value of AMOUNT . Then order this summed up column within ID groups and return the SUB_ID values of the maximum value.

We can get the summation by:

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

This will return a Series with MultiIndex

 ID    SUB_ID    

 101     1        80
         2        50
 102     3        20
         4        20

From here I would like to return [1,3] ([1,4] is also accepted as the two values in the 102 group are the same, but we want to return only one value per group!)

Obviously we can loop and pick the max but I am trying to find out the most efficient way possible. This operation will be applied to millions of rows.

Answer 1

This is one way. As your dataset is large, I strongly recommend you avoid lambda functions since these are not applied in a vectorised fashion.

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

---

Example

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

I think you can use nlargest :

df1.groupby('ID').nlargest(1).index.get_level_values(level='SUB_ID').tolist()

# [1, 3]