Включать только первые N вхождений вторичного столбца в pandas groupby

Question

Я пытаюсь сделать в pandas то, что сделал бы с коррелированным подзапросом в SQL, и не смог найти не итеративное и / или более быстрое решение. Рассмотрим (гораздо больший) фрейм данных игроков, игр и статистики для этих игр, который можно определить следующим образом:

import pandas as pd

def make_mock_df():
    mock_data = {
        "player": ["Bob", "Bob", "Bob", "Bob", "Bob", "Jim", "Jim", "Jim", "Jim", "Jim"],
        "game_id": ["red", "red", "green", "green", "blue", "red", "yellow", "yellow", "blue", "blue"],
        "stat_id": [1, 2, 1, 2, 1, 1, 1, 2, 1, 2],
        "value": [1, 3, 2, 4, 3, 1, 1, 2, 4, 2],
    }
    return pd.DataFrame(mock_data)

На самом деле идентификатор игры будет UUID, и данные уже будут сортировать сначала по игроку, а затем по дате игры, недавно-сначала.

Я хочу создать информационный фрейм, который будет усреднять каждый тип статистики за последние N игр на игрока. Это становится непросто, поскольку не каждый игрок играет в каждую игру (по факту насмешек в данных), и поэтому я сделал бы соответствующий подзапрос в SQL. До сих пор это итеративное решение, которое разбивает информационный кадр по игроку, затем объединяет каждый ограниченный информационный кадр игрока, чтобы перестроить целое:

def filter_last_n_games(df, n_games):
    last_n_games_per_player = pd.DataFrame()

    for _, player_group in df.groupby("player"):
        # this works because games are sorted already
        # and appear sorted in the unique() series as well
        last_n_games = player_group.game_id.unique()[0:n_games]
        # can't do nlargest because games have string IDs
        last_n_mask = player_group.game_id.isin(last_n_games)
        player_last_n = player_group[last_n_mask]
        # build up last N games df piece by piece...
        last_n_games_per_player = pd.concat(
            [last_n_games_per_player, player_last_n]
        )

    return last_n_games_per_player

def avg_stats_last_n_games(df, n_games):
    last_n_games_per_player = filter_last_n_games(df, n_games)

    # calculate averages per game (easy part)
    last_n_avgs = (
        last_n_games_per_player
        .groupby(["player", "stat_id"])
        .mean()
        .filter(["player", "stats_id", "value"])
        .rename(columns={"value": "avg_value"})
    )

    return last_n_avgs

И приводит к выводу как таковому:

stats_df = make_mock_df()
avg_stats_last_n_games(stats_df, 2)
>>>                 avg_value
>>> player stat_id           
>>> Bob    1              1.5
>>>        2              3.5
>>> Jim    1              1.0
>>>        2              2.0

Есть ли лучшая замена функции filter_last_n_games(), которая бы в полной мере воспользовалась возможностями pandas 'и позволила бы мне отказаться от for-l oop?

Answer 1

IIU C, эквивалентный выполнению подзапроса:

(df[df.groupby(['player','stat_id']).cumcount() < 2] #filter dataframe to first two records
   .groupby(['player','stat_id'])['value'].mean())  #then groupby and take average

Выход:

player  stat_id
Bob     1          1.5
        2          3.5
Jim     1          1.0
        2          2.0
Name: value, dtype: float64

Answer 2

Для каждой группы, для выбора последних n игр требуется,

1. всего # элементов в группе tcount.
2. тегирование элементов в индексе группы icount.

Для фиктивного кадра данных:

# With your `make_mock_df()` I am getting

#   player game_id  stat_id  value
# 0    Bob     red        1      1
# 1    Bob     red        2      3
# 2    Bob   green        1      2
# 3    Bob   green        2      4
# 4    Bob    blue        1      3
# 5    Jim     red        1      1
# 6    Jim  yellow        1      1
# 7    Jim  yellow        2      2
# 8    Jim    blue        1      4
# 9    Jim    blue        2      2

def add_counters(grp):
    grp['tcount'] = grp.shape[0]
    grp['icount'] = np.arange(grp.shape[0]) + 1
    return grp


gdf = make_mock_df()
gdf = gdf.groupby(['player', 'game_id']).apply(lambda x: add_counters(x))

# Selecting last `N` games
N=2;

gdf = gdf[(gdf['tcount'] >= N) & (gdf['icount'] <= N)]

# After this filtering step i Get
#   player game_id  stat_id  value  tcount  icount
# 0    Bob     red        1      1       2       1
# 1    Bob     red        2      3       2       2
# 2    Bob   green        1      2       2       1
# 3    Bob   green        2      4       2       2
# 6    Jim  yellow        1      1       2       1
# 7    Jim  yellow        2      2       2       2
# 8    Jim    blue        1      4       2       1
# 9    Jim    blue        2      2       2       2


gdf.groupby(['player', 'stat_id'])[['value']].agg('mean').rename(columns={'value': 'avg_value'})

#                 avg_value
# player stat_id
# Bob    1              1.5
#        2              3.5
# Jim    1              2.5
#        2              2.0

Проверьте правки выше