选择每个组的最大值

时间:2018-10-03 04:08:07

标签: python pandas

所以我有一个包含多个列和一个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() 

任何人都对我如何获得期望的结果提出建议

1 个答案:

答案 0 :(得分:2)

设置

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

另一种解决方案是使用meltgroupby

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