Pandas DataFrame n-самый большой, который зависит от группы - PullRequest
Новая
       6

Pandas DataFrame n-самый большой, который зависит от группы

2 голосов
/ 21 апреля 2020

Я хочу взять n наибольшие значения DataFrame для каждого Date, но я хочу, чтобы n варьировалось по дате. Так, например, этот limits DataFrame указывает количество значений, которые я хочу для каждой даты: 

np.random.seed(456)
limits = pd.DataFrame(np.random.randint(2,5,5), pd.date_range('2020-01-01','2020-01-05').tolist(), columns=['limit'])

            limit
2020-01-01      4
2020-01-02      2
2020-01-03      4
2020-01-04      2
2020-01-05      3

А вот пример DataFrame, к которому я хочу применить те: 

j = [(a, b) for a in ['A','B','C','D','E'] for b in pd.date_range('2020-01-01','2020-01-05').tolist()]
i = pd.MultiIndex.from_tuples(j, names=['Name','Date'])
df = pd.DataFrame(np.random.randn(25), i, columns=['Vals'])

                     Vals
Name Date                
A    2020-01-01 -1.240210
     2020-01-02 -0.954311
     2020-01-03 -0.468707
     2020-01-04 -0.861229
     2020-01-05  0.138360
B    2020-01-01 -0.164922
     2020-01-02 -0.257626
     2020-01-03 -1.200235
...

С постоянной ссылкой n = 2 я могу получить 2 самых больших значения, используя:

df.groupby(['Date']).apply(lambda x: (x.sort_values('Vals').head(2))).reset_index(level=0, drop=True)

                     Vals
Name Date                
A    2020-01-01 -1.240210
E    2020-01-01 -1.095603
D    2020-01-02 -1.298098
A    2020-01-02 -0.954311
...

Но как мне получить число строк, указанное limits на каждую дату?

Ответы [ 3 ]

3 голосов
/ 21 апреля 2020

Вы можете назначить limit в качестве нового столбца, а затем использовать query для фильтрации: 

(df.assign(limit=limits.loc[df.index.get_level_values('Date'),'limit'].values,
          order=lambda x: x.sort_values('Vals', ascending=False).groupby('Date').cumcount()          # sort is needed for `nlargest`
         )
   .query('order< limit')
   .drop(['order','limit'], axis=1)
)

Вывод: 

                     Vals
Name Date                
A    2020-01-01  1.246749
     2020-01-02 -0.079275
     2020-01-03 -0.636896
     2020-01-04  0.013802
     2020-01-05 -1.397262
B    2020-01-01  1.726135
     2020-01-02 -0.491877
     2020-01-03  0.254206
     2020-01-04 -0.268168
     2020-01-05 -0.066552
C    2020-01-01 -1.017655
     2020-01-03  0.671070
     2020-01-05 -0.135537
D    2020-01-01  1.813671
     2020-01-03 -0.882443
2 голосов
/ 21 апреля 2020

pd.concat

pd.concat([
    d.nlargest(limits.limit[date], columns=['Vals'])
    for date, d in df.groupby('Date')
])

Использовать маску

При этом используется та же техника (с использованием cumcount), что и Quang Hoang 

d = df.sort_values(['Date', 'Vals'], ascending=[True, False])
c = d.groupby('Date').cumcount() + 1
d[c <= d.index.get_level_values(1).map(limits.limit)]

                     Vals
Name Date                
A    2020-01-01  1.350509
D    2020-01-01  1.157552
E    2020-01-01  1.139873
C    2020-01-02  1.944702
A    2020-01-02  1.629589
E    2020-01-02  0.136372
C    2020-01-03  1.915676
A    2020-01-03  0.301966
D    2020-01-03 -0.088752
E    2020-01-03 -0.366948
C    2020-01-04  0.920348
A    2020-01-04  0.449483
C    2020-01-05  0.936398
B    2020-01-05  0.237851
E    2020-01-05  0.107640
A    2020-01-05 -0.345811
​
2 голосов
/ 21 апреля 2020

Лямбда-функция содержит индекс для каждой строки, которую она обрабатывает, и доступ к ней можно получить с помощью ее свойства .name. Итак, с limits и df, как определено в вопросе: 

df.groupby(['Date']).apply(lambda x:
     (x.sort_values('Vals').head(limits.loc[x.name].limit)))
  .reset_index(level=0, drop=True)

возвращает то, что вы хотели: 

                     Vals
Name Date                
A    2020-01-01 -1.240210
E    2020-01-01 -1.095603
C    2020-01-01 -0.510581
B    2020-01-01 -0.164922
D    2020-01-02 -1.298098
A    2020-01-02 -0.954311
B    2020-01-03 -1.200235
...
...