如何高效基于条件从DataFrame其他列生成新列？

问题背景 #

给定如下结构的DataFrame（df，省略号表示更多列）：

Type  Price1  Price2  Price3  Price4  Price5  ... ... 
A       nan     1       nan     nan     2           
A       nan     3       nan     nan     2
B       nan     nan     4       5       nan
B       nan     nan     6       7       nan
C       nan     2       nan     nan     1  
C       nan     4       nan     nan     3
D       1       8       nan     nan     nan
D       9       6       nan     nan     nan

需要转换为如下结构：

Type   newcol1    newcol2  ...  ...
     A      2           1
     A      2           3
     B      4           5
     B      6           7
     C      1           2   
     C      3           4   
     D      1           8       
     D      9           6

列值映射规则：

• 当Type为'A'时：Price5 → newcol1，Price2 → newcol2
• 当Type为'B'时：Price3 → newcol1，Price4 → newcol2
• 当Type为'C'时：Price5 → newcol1，Price2 → newcol2
• 当Type为'D'时：Price1 → newcol1，Price2 → newcol2

原方法使用mask函数逐个处理每种类型，步骤繁琐，需要更高效的实现方式。

高效实现方案 #

方法1：映射字典+lookup向量化生成 #

先定义类型到源列的映射字典，后续新增类型只需修改此字典：

col_mapping = {
    'A': ('Price5', 'Price2'),
    'B': ('Price3', 'Price4'),
    'C': ('Price5', 'Price2'),
    'D': ('Price1', 'Price2')
}

通过map获取每行对应的源列名，再用lookup批量提取值生成新列：

# 为每行匹配newcol1/newcol2对应的源列名
newcol1_src = df['Type'].map(lambda x: col_mapping[x][0])
newcol2_src = df['Type'].map(lambda x: col_mapping[x][1])

# 批量提取对应列的值
df['newcol1'] = df.lookup(df.index, newcol1_src)
df['newcol2'] = df.lookup(df.index, newcol2_src)

# 若只需保留Type、新列及原其他非价格列，可筛选：
# keep_cols = ['Type', 'newcol1', 'newcol2'] + [col for col in df.columns if col not in ['Price1','Price2','Price3','Price4','Price5']]
# df = df[keep_cols]

方法2：np.select批量赋值 #

如果偏好使用条件判断的方式，可借助np.select实现：

import numpy as np

# 定义匹配条件
conditions = [
    df['Type'] == 'A',
    df['Type'] == 'B',
    df['Type'] == 'C',
    df['Type'] == 'D'
]

# 对应每个条件的newcol1取值列
newcol1_values = [
    df['Price5'],
    df['Price3'],
    df['Price5'],
    df['Price1']
]

# 对应每个条件的newcol2取值列
newcol2_values = [
    df['Price2'],
    df['Price4'],
    df['Price2'],
    df['Price2']
]

# 生成新列
df['newcol1'] = np.select(conditions, newcol1_values)
df['newcol2'] = np.select(conditions, newcol2_values)

方法3：分组处理（适合复杂映射逻辑） #

如果后续映射规则涉及更复杂的行内计算，可按Type分组后单独处理：

def process_group(group):
    type_val = group['Type'].iloc[0]
    col1, col2 = col_mapping[type_val]
    group['newcol1'] = group[col1]
    group['newcol2'] = group[col2]
    return group

df = df.groupby('Type').apply(process_group)

以上方法均无需单独处理每个类型，代码复用性和扩展性更强，新增类型时仅需更新映射字典或条件列表即可。

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