如何在Pandas DataFrame中将首次出现ONE-OFF后的列设为NaN

我们需要处理Pandas DataFrame：对每行来说，找到第一个值为"ONE-OFF"的列，把该列之后的所有列值替换为NaN；如果某行没有"ONE-OFF"，则保持该行不变。

步骤1：导入必要库 #

import pandas as pd
import numpy as np

步骤2：核心处理函数 #

def mask_after_first_one_off(df):
    # 获取每行第一个"ONE-OFF"的列索引，无匹配则返回第一个列索引（后续掩码会自动跳过该行）
    first_off_pos = df.eq('ONE-OFF').idxmax(axis=1)
    # 生成与DataFrame形状一致的列索引矩阵，用于逐行位置比较
    col_indices = np.tile(df.columns, (len(df), 1))
    # 生成掩码：标记每行中在第一个"ONE-OFF"之后的列
    mask = col_indices > first_off_pos.values[:, np.newaxis]
    # 复制原数据避免修改原始输入
    result_df = df.copy()
    # 应用掩码将对应位置设为NaN
    result_df[mask] = np.nan
    return result_df

步骤3：测试示例输入 #

# 示例输入
input_df = pd.DataFrame(
    {
         1: {'15': 'Normal'},
         2: {'15': 'Normal'},
         3: {'15': 'Normal'},
         4: {'15': 'ONE-OFF'},
         5: {'15': 'Normal'},
         6: {'15': 'Normal'},
    }
)

# 处理得到结果
output_df = mask_after_first_one_off(input_df)
print(output_df)

运行后输出与示例一致：

1       2       3       4    5    6
15  Normal  Normal  Normal  ONE-OFF  NaN  NaN

代码逻辑说明 #

• df.eq('ONE-OFF').idxmax(axis=1)：逐行判断值是否为"ONE-OFF"，返回每行第一个匹配项的列索引；无匹配时返回首列索引，后续比较会确保该行无修改。
• col_indices：将列索引重复为与DataFrame行数一致的矩阵，实现逐行的列位置对比。
• 布尔掩码mask：标记所有需要替换为NaN的位置，最后通过掩码批量赋值完成修改。

边界情况验证 #

如果存在无"ONE-OFF"的行，函数会自动保留该行原始数据：

test_df = pd.DataFrame(
    {
         1: {'10': 'Normal', '15': 'Normal'},
         2: {'10': 'Normal', '15': 'ONE-OFF'},
         3: {'10': 'Normal', '15': 'Normal'},
         4: {'10': 'Normal', '15': 'Normal'},
    }
)

result_df = mask_after_first_one_off(test_df)
print(result_df)

输出：

1       2       3       4
10  Normal  Normal  Normal  Normal
15  Normal  ONE-OFF     NaN     NaN

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