获取最大值行的代码优化idxmax()

53 阅读1分钟

这是我写的获取每个itemid下的最大值行,需要3步完成

import pandas as pd

# 示例数据框
price_df= pd.DataFrame({
    'itemId': ["A", "A", "A", "B", "B","B"," C"," C"," C"],
    'discountePrice': [1, 1, 1, 2, 3, 1, 4,5,5]
})
display(price_df)
# 得到itemid的最大值
price_df_drop_only_max=price_df.groupby("itemId").max()
# itemid的最大值与原表链接
price_df=price_df.merge(price_df_drop_only_max,on='itemId',how='left',suffixes=('','_max'))
#筛选出所有最大值行
price_df['max']=price_df['discountePrice_max']==price_df['discountePrice']

优化方案:

import pandas as pd

# 示例数据框
price_df = pd.DataFrame({
    'itemId': ["A", "A", "A", "B", "B", "B", "C", "C", "C"],
    'discountePrice': [1, 1, 1, 2, 3, 1, 4, 5, 5]
})
display(price_df)

# 得到每个 itemId 的最大 discountePrice 行
price_df_drop_only_max = price_df.loc[price_df.groupby('itemId')['discountePrice'].idxmax()]

display(price_df_drop_only_max)

解释:

  1. 创建示例数据框:包含 itemIddiscountePrice 列。
  2. 使用 groupbyidxmax 筛选最大值行
    • price_df_drop_1.groupby('itemId')['discountePrice'].idxmax() 计算每个 itemId 组内 discountePrice 的最大值的索引。
    • price_df_drop_1.loc[...] 根据索引筛选出最大值行。

这个方法更简洁,并且直接返回每个 itemId 的最大 discountePrice 行。

运行结果:

price_df_drop_only_max 将包含每个 itemId 的最大 discountePrice 行:

  itemId  discountePrice
2      A               1
4      B               3
7      C               5

这个方法不仅简化了代码,而且更加高效地实现了目标。