在Pandas DataFrame中按分组取最大值并保留其他列

背景 #

免责声明：这是我第一次使用Pandas。

我有一个无分层索引的DataFrame，结构如下：

我希望获取每个州的最大百分比，以及该百分比对应的政党，预期输出示例如下：

尝试的代码 #

# 提取需要的列
df_h = pd.DataFrame(house_polls, columns=['state','party','pct'])
df_s = pd.DataFrame(senate_polls, columns=['state','party','pct'])

# 按州和政党分组计算平均百分比，并排序
cols = ["state", "party"]
df_h_avg = df_h.groupby(cols, as_index=False).mean(numeric_only=True).sort_values(by=cols)
df_s_avg = df_s.groupby(cols, as_index=False).mean(numeric_only=True).sort_values(by=cols)

# 获取每个州的最大百分比
cols1 = ["state"]
house_results = df_h_avg.groupby(by=cols1).max(numeric_only=True)
senate_results = df_s_avg.groupby(by=cols1).max(numeric_only=True)

# 怎么关联对应的政党？

遇到的问题 #

运行上述代码后，仅得到每个州的最大百分比，丢失了对应的政党信息，结果如下：

解决方案 #

问题出在groupby().max(numeric_only=True)仅保留数值列的最大值，无法关联对应的政党信息。以下是几种可行的解决方法：

方法1：用idxmax()定位最大值行索引 #

先找到每个州百分比最大值对应的行索引，再通过索引提取完整行数据：

# 处理众议院数据
house_max_idx = df_h_avg.groupby('state')['pct'].idxmax()
house_results = df_h_avg.loc[house_max_idx].reset_index(drop=True)

# 处理参议院数据
senate_max_idx = df_s_avg.groupby('state')['pct'].idxmax()
senate_results = df_s_avg.loc[senate_max_idx].reset_index(drop=True)

方法2：排序后去重 #

先按州升序、百分比降序排序，再保留每个州的第一行（即百分比最高的行）：

# 处理众议院数据
house_results = df_h_avg.sort_values(['state', 'pct'], ascending=[True, False])\
                        .drop_duplicates('state').reset_index(drop=True)

# 处理参议院数据
senate_results = df_s_avg.sort_values(['state', 'pct'], ascending=[True, False])\
                          .drop_duplicates('state').reset_index(drop=True)

方法3：自定义聚合函数 #

通过apply()实现自定义聚合，同时获取最大百分比及对应政党：

def get_max_party(group):
    max_row = group[group['pct'] == group['pct'].max()]
    return max_row.iloc[0] if not max_row.empty else None

# 处理众议院数据
house_results = df_h_avg.groupby('state').apply(get_max_party).reset_index(drop=True)

# 处理参议院数据
senate_results = df_s_avg.groupby('state').apply(get_max_party).reset_index(drop=True)

以上方法均可保留每个州百分比最高的政党及对应数值，匹配预期输出。

内容的提问来源于stack exchange，提问作者lordcarrot