Python Pandas：查找首辆新车列索引并关联对应字段的问题

作为Python新手（习惯SAS逻辑），在Pandas DataFrame处理中遇到两个核心问题：

1. 在NU1-NU5列中查找第一个值为"N"（新车）的列时，无新车的行（如测试数据第3行全为"U"）会返回错误的列名，而非空值/标识值。
2. 无法高效利用找到的首辆新车位置，关联获取Car、purchase_date、price等同结构字段的对应值。

测试数据与初始代码如下：

import pandas as pd
import numpy as np

# 创建测试DataFrame
df = pd.DataFrame([[12345,"Name1","Name1","Name2","","","A","U","N","N","",""], 
               [45678,"Name3","Name1","Name2","","","S","N","N","N","",""],
               [45679,"Name2","Name2","Name2","","","S","U","U","U","",""],
               [98765,"Name3","Name2","Name2","","","","CPO","U","N","",""]], 
           columns=("ID","Car1","Car2","Car3","Car4","Car5","Income",
                    "NU1","NU2","NU3","NU4","NU5") )

nu_array = df[["NU1","NU2","NU3","NU4","NU5"]] 
nameplate_array = df[["Car1","Car2","Car3","Car4","Car5"]] 

# 查找第一个新车列名（存在问题：无新车行仍返回错误列名）
df['first_nu_col_name'] = df.eq("N").T.idxmax()  
df['col_num'] = df['first_nu_col_name'].str[2:]
df['first_nu_veh_col_name'] = "Car" + df['col_num']

1. 修正无新车行的错误返回值 #

原代码中idxmax()在全行无"N"时会返回第一个列名（因为布尔值全为False，idxmax取第一个索引）。我们可以先标记每行是否存在新车，再对无新车的行设置空值：

# 仅针对NU1-NU5列判断是否为"N"
nu_cols = ["NU1","NU2","NU3","NU4","NU5"]
nu_bool = df[nu_cols].eq("N")

# 获取第一个"N"对应的列名，无则设为NaN
df['first_nu_col_name'] = nu_bool.idxmax(axis=1)
# 对无新车的行，将列名置为NaN
df.loc[~nu_bool.any(axis=1), 'first_nu_col_name'] = np.nan

这样第3行的first_nu_col_name会变成NaN，后续的col_num和first_nu_veh_col_name也会自动变为NaN，避免错误列名。

2. 高效关联同结构字段 #

不需要通过拼接列名的方式取值，直接利用位置索引更高效且符合Python风格：

步骤1：获取第一个新车的位置索引 #

# 获取每行第一个"N"的列索引（从0开始，对应NU1-NU5的顺序）
df['first_nu_pos'] = nu_bool.idxmax(axis=1).str.extract(r'(\d+)').astype(float) - 1
# 无新车的行，位置设为NaN
df.loc[~nu_bool.any(axis=1), 'first_nu_pos'] = np.nan

步骤2：批量获取对应字段值 #

以Car列为例，直接用位置索引提取对应值：

car_cols = ["Car1","Car2","Car3","Car4","Car5"]
# 用apply按行提取对应位置的Car值
df['first_nu_car'] = df.apply(lambda row: row[car_cols][int(row['first_nu_pos'])] if pd.notna(row['first_nu_pos']) else np.nan, axis=1)

如果需要扩展到purchase_date1-5、price1-5等同结构字段，只需替换car_cols为对应字段列表即可：

# 示例：获取对应purchase_date值
date_cols = ["purchase_date1","purchase_date2","purchase_date3","purchase_date4","purchase_date5"]
df['first_nu_date'] = df.apply(lambda row: row[date_cols][int(row['first_nu_pos'])] if pd.notna(row['first_nu_pos']) else np.nan, axis=1)

# 示例：获取对应price值
price_cols = ["price1","price2","price3","price4","price5"]
df['first_nu_price'] = df.apply(lambda row: row[price_cols][int(row['first_nu_pos'])] if pd.notna(row['first_nu_pos']) else np.nan, axis=1)

完整代码 #

import pandas as pd
import numpy as np

# 创建测试DataFrame
df = pd.DataFrame([[12345,"Name1","Name1","Name2","","","A","U","N","N","",""], 
               [45678,"Name3","Name1","Name2","","","S","N","N","N","",""],
               [45679,"Name2","Name2","Name2","","","S","U","U","U","",""],
               [98765,"Name3","Name2","Name2","","","","CPO","U","N","",""]], 
           columns=("ID","Car1","Car2","Car3","Car4","Car5","Income",
                    "NU1","NU2","NU3","NU4","NU5") )

# 定义字段组
nu_cols = ["NU1","NU2","NU3","NU4","NU5"]
car_cols = ["Car1","Car2","Car3","Car4","Car5"]

# 生成NU列的布尔矩阵
nu_bool = df[nu_cols].eq("N")

# 1. 修正无新车行的错误列名
df['first_nu_col_name'] = nu_bool.idxmax(axis=1)
df.loc[~nu_bool.any(axis=1), 'first_nu_col_name'] = np.nan

# 2. 获取第一个新车的位置索引（从0开始）
df['first_nu_pos'] = nu_bool.idxmax(axis=1).str.extract(r'(\d+)').astype(float) - 1
df.loc[~nu_bool.any(axis=1), 'first_nu_pos'] = np.nan

# 3. 提取对应Car值
df['first_nu_car'] = df.apply(lambda row: row[car_cols][int(row['first_nu_pos'])] if pd.notna(row['first_nu_pos']) else np.nan, axis=1)

# 扩展示例：若存在purchase_date和price字段，提取对应值
# 假设df已包含这些字段
# date_cols = ["purchase_date1","purchase_date2","purchase_date3","purchase_date4","purchase_date5"]
# df['first_nu_date'] = df.apply(lambda row: row[date_cols][int(row['first_nu_pos'])] if pd.notna(row['first_nu_pos']) else np.nan, axis=1)
# price_cols = ["price1","price2","price3","price4","price5"]
# df['first_nu_price'] = df.apply(lambda row: row[price_cols][int(row['first_nu_pos'])] if pd.notna(row['first_nu_pos']) else np.nan, axis=1)

print(df[['ID', 'first_nu_col_name', 'first_nu_pos', 'first_nu_car']])

运行结果 #

ID first_nu_col_name  first_nu_pos first_nu_car
0  12345               NU2           1.0        Name1
1  45678               NU1           0.0        Name3
2  45679                NaN           NaN          NaN
3  98765               NU3           2.0        Name2

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