基于另一列的值将DataFrame指定列值替换为NaN的实现方法

原始DataFrame #

import pandas as pd
import numpy as np

df = pd.DataFrame(
    {
        'a': [np.nan, np.nan, np.nan, 3333, np.nan,  np.nan, 10, np.nan, np.nan, np.nan, np.nan, 200, 100],
        'b': [np.nan, 20, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, 100, np.nan, np.nan, np.nan, np.nan]
    }
)

期望输出 #

a      b
0      NaN    NaN
1      NaN   20.0
2      NaN    NaN
3   3333.0    NaN
4      NaN    NaN
5      NaN    NaN
6      NaN    NaN
7      NaN    NaN
8      NaN  100.0
9      NaN    NaN
10     NaN    NaN
11   200.0    NaN
12     NaN    NaN

规则说明 #

当b列出现非NaN值时，仅保留该位置之后a列的第一个非NaN值，后续所有a列的非NaN值需替换为NaN，直到b列出现下一个非NaN值。例如：

• b列出现20后，保留a列的3333，将后续的10替换为NaN
• b列出现100后，保留a列的200，将后续的100替换为NaN

解决方案 #

通过分组标记的方式实现需求，具体代码如下：

# 生成分组标签：b列每出现一个非NaN值，分组号递增，其余行继承前一个分组号
df['group'] = df['b'].notna().cumsum()

# 对每个分组，标记a列中第一个非NaN值的位置
mask = df.groupby('group')['a'].transform(lambda x: x.notna().cumsum() == 1)

# 仅保留mask为True的a值，其余置为NaN
df['a'] = df['a'].where(mask, np.nan)

# 删除临时分组列
df = df.drop('group', axis=1)

print(df)

运行上述代码后，即可得到符合要求的输出结果。

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