python - 选择每组的最大值

Question

所以我有一个包含多个列和一个 id 列的 pandas 数据框。

df = pd.DataFrame(np.random.randn(6,4), columns=list('ABCD'))
df['id'] = ['CA', 'CA', 'CA', 'FL', 'FL', 'FL']
df['technique'] = ['one', 'two', 'three', 'one', 'two', 'three']
df

我想按 id 列进行分组并选择概率最高的行。所以它可能看起来像这样。

id   highest_prob   technique
CA   B               three 
FL   C               one

我尝试过类似的事情，但这只能让我成功一半。

df.groupby('id', as_index=False)[['A','B','C','D']].max() 

任何人都可以建议我如何获得所需的结果

Answer 1

设置

np.random.seed(0)  # Add seed to reproduce results. 
df = pd.DataFrame(np.random.randn(6,4), columns=list('ABCD'))
df['id'] = ['CA', 'CA', 'CA', 'FL', 'FL', 'FL']
df['technique'] = ['one', 'two', 'three', 'one', 'two', 'three']

您可以melt，使用sort_values排序，并使用drop_duplicates删除重复项:

(df.melt(['id', 'technique'])
   .sort_values(['id', 'value'], ascending=[True, False])
   .drop_duplicates('id')
   .drop('value', 1)
   .reset_index(drop=True)
   .rename({'variable': 'highest_prob'}, axis=1))

   id technique highest_prob
0  CA       one            D
1  FL       two            A

<小时/>

另一种解决方案是使用melt和groupby:

v = df.melt(['id', 'technique'])
(v.iloc[v.groupby('id').value.idxmax()]
  .drop('value', 1)
  .reset_index(drop=True)
  .rename({'variable': 'highest_prob'}, axis=1))

   id technique highest_prob
0  CA       one            D
1  FL       two            A