基于行内条件从多组对应列选取值生成新列的优化方法问询

原始数据 #

期望结果 #

需要新增列y，根据Duration的取值，从对应的y_dur1/y_dur2/y_dur3中提取值，结果如下：

实际场景需处理5组共30列（每组6列），当前使用的代码如下：

df = pd.read_csv(path)

cols = df.columns
dur_cols = pd.Series([col for col in df.columns if "_Dur" in col])

val_cols = dur_cols.replace('\\d',"", regex=True)

val_cols.drop_duplicates(inplace=True)

for col in val_cols:
    df[col] = np.where(df["Duration_to_use"]=="dur1",df[col+"1"],np.where(df["Duration_to_use"]=="dur2",df[col+"2"],df[col+"3"]))

优化方案 #

方案1：使用lookup方法（最简洁高效） #

pandas的lookup方法可直接按行匹配列名提取值，单组场景一行搞定：

df['y'] = df.lookup(df.index, 'y_' + df['Duration'])

多组场景（比如前缀为y_、x_、z_等）可批量处理：

# 提取所有组前缀（假设列名格式为"前缀_durN"）
prefixes = {col.split('_')[0] for col in df.columns if '_dur' in col.lower()}

for prefix in prefixes:
    df[prefix] = df.lookup(df.index, f'{prefix}_' + df['Duration'])

方案2：使用apply行操作（逻辑直观） #

逐行处理逻辑清晰，适合小数据集：

df['y'] = df.apply(lambda row: row[f'y_{row["Duration"]}'], axis=1)

多组场景批量处理：

prefixes = {col.split('_')[0] for col in df.columns if '_dur' in col.lower()}
for prefix in prefixes:
    df[prefix] = df.apply(lambda row: row[f'{prefix}_{row["Duration"]}'], axis=1)

方案3：使用melt+merge（通用性强） #

如果后续有复杂匹配需求，可先转长表再合并：

# 转长表并拆分列名前缀与Duration标识
melted = df.melt(id_vars=['Name', 'Duration'], 
                 var_name='col', 
                 value_name='value')
melted[['prefix', 'dur']] = melted['col'].str.split('_', expand=True)

# 匹配Duration对应的行
matched = melted[melted['dur'] == melted['Duration']]

# 合并回原表并恢复宽表格式
df = df.merge(matched[['Name', 'prefix', 'value']], 
              on='Name', 
              how='left')
df = df.pivot(index=['Name', 'Duration', 'y_dur1', 'y_dur2', 'y_dur3'], 
              columns='prefix', 
              values='value').reset_index().rename_axis(None, axis=1)

方案对比 #

• lookup：效率最高、代码最简洁，是规则化列名匹配场景的最优选择。
• apply：逻辑直观但效率较低，数据量大时不推荐。
• melt+merge：通用性强，适合列名规则不统一或需额外处理的场景，但代码稍繁琐。

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