используйте group by, чтобы получить n наименьших значений, но с дубликатами - PullRequest
2 голосов
/ 06 августа 2020

Предположим, у меня есть pandas DataFrame вроде этого:

>>> df = pd.DataFrame({'id':[1,1,1,1,1,2,2,2,2,2,2,3,4],'value':[1,1,1,1,3,1,2,2,3,3,4,1,1]})
>>> df
id  value
1      1
1      1
1      1
1      1
1      3
2      1
2      2
2      2
2      3
2      3
2      4
3      1
4      1

Я хочу получить новый DataFrame с верхними 2 (ну, действительно, n значениями) значениями для каждого идентификатора включая дубликаты , например:

   id  value
0   1      1
1   1      1
3   1      1
4   1      1
5   1      3
6   2      1
7   2      2
8   2      2
9   3      1
10  4      1

Я пробовал использовать head () и nsmallest (), но я думаю, что они не будут включать дубликаты. Есть ли лучший способ сделать это?

Отредактировано, чтобы было понятно, что мне нужно более 2 записей на группу, если имеется более 2 дубликатов

1 Ответ

2 голосов
/ 06 августа 2020

Используйте DataFrame.drop_duplicates на первом этапе, затем получите максимальные значения и используйте последний раз DataFrame.merge:

df1 = df.drop_duplicates(['id','value']).sort_values(['id','value']).groupby('id').head(2)
df = df.merge(df1)
print (df)
   id  value
0   1      1
1   1      1
2   1      2
3   1      2
4   2      1
5   2      2
6   2      2
7   3      1
8   4      1
df = pd.DataFrame({'id':[1,1,1,1,1,2,2,2,2,2,2,3,4],'value':[1,1,1,1,3,1,2,2,3,3,4,1,1]})
    
df1 = df.drop_duplicates(['id','value']).sort_values(['id','value']).groupby('id').head(2)
df = df.merge(df1)
print (df)
   id  value
0   1      1
1   1      1
2   1      1
3   1      1
4   1      3
5   2      1
6   2      2
7   2      2
8   3      1
9   4      1

Или используйте пользовательскую лямбда-функцию с GroupBy.transform и фильтруйте в boolean indexing:

df = df[df.groupby('id')['value'].transform(lambda x: x.isin(sorted(set(x))[:2]))]
print (df)
    id  value
0    1      1
1    1      1
2    1      2
3    1      2
5    2      1
6    2      2
7    2      2
11   3      1
12   4      1
df = df[df.groupby('id')['value'].transform(lambda x: x.isin(sorted(set(x))[:2]))]
print (df)
    id  value
0    1      1
1    1      1
2    1      1
3    1      1
4    1      3
5    2      1
6    2      2
7    2      2
11   3      1
12   4      1
...