如何将DataFrame分组后每组最后n个非NaN值设为NaN？

我有一个包含若干分组变量和数值变量的DataFrame，目标是将每组中的最后n个非NaN值设置为NaN。

示例数据 #

import pandas as pd
import numpy as np

df = pd.DataFrame({'id':[1,1,1,2,2],
                  'value':[1,2,np.nan, 9,8]})

原始输出 #

id  value
0   1    1.0
1   1    2.0
2   1    NaN
3   2    9.0
4   2    8.0

期望结果（当n=1时） #

id  value
0   1    1.0
1   1    NaN
2   1    NaN
3   2    9.0
4   2    NaN

解决方案 #

可以通过groupby结合反向累计计数实现：

n = 1

# 对每个分组，筛选出需要设为NaN的最后n个非空值位置
mask = df.groupby('id')['value'].apply(
    lambda x: x.notna() & (x.notna()[::-1].cumsum() <= n)
)

# 将标记位置的value设为NaN
df.loc[mask, 'value'] = np.nan

代码说明 #

• x.notna()：筛选当前分组中数值非空的位置
• x.notna()[::-1].cumsum()：反转非空位置序列后累计计数，得到从最后一个非NaN值开始的序号
• <=n：锁定每组最后n个非NaN值的位置
• 用loc定位目标位置并设置为NaN

内容的提问来源于stack exchange，提问作者TiTo