pandas で グループごとに ffill() と bfill() を両方適用したい

2020-10-19 13:16 修正。新しい列を作る必要は特にない。

pandas を使うと、例えば

horses.loc[:, 'horse_weight'] = horses.groupby('horse_id')['horse_weight'].ffill()

とすることで、horse_weight の中で欠損した値を、horse_id でグループ化した中での1つ前の値を使って埋めることができるが、1つ前の値が存在しない場合に、1つ後の値を使って補完 (bfill()) したい場合はどうしたら良いか。

# これは正しくない
horses.loc[:, 'horse_weight'] = horses.groupby('horse_id')['horse_weight'].ffill().bfill()

これはいけない。なぜなら ffill() が返す型は SeriesGroupBy ではなく Series であり、グループ化が解除されているため、グループに関係なく 1 つ後の値を使ってしまう。

Stack Overflow に回答があるが、apply() を使うのが正しい。

horses.loc[:, 'horse_weight'] = horses.groupby('horse_id')['horse_weight'].apply(lambda x: x.ffill().bfill())

ただ、実際やってみるとこれは結構遅かった。たぶん普通に SeriesGroupBy.ffill() やるときは何か良い感じの最適化がかかってる一方で、apply() の場合は律儀にグループ数分のループが走るからだと思う（未検証）。

ということで、愚直に

horses.loc[:, 'horse_weight_filled'] = horses.groupby('horse_id')['horse_weight'].ffill() # 別の列に ffill() した結果を入れておく
horses.loc[:, 'horse_weight_filled'] = horses.groupby('horse_id')['horse_weight_filled'].bfill()

よく考えたら別に新しい列を作る必要はなかった。

horses.loc[:, 'horse_weight'] = horses.groupby('horse_id')['horse_weight'].ffill()
horses.loc[:, 'horse_weight'] = horses.groupby('horse_id')['horse_weight'].bfill()

で良い。