重塑DataFrame时移除NA值并创建新列的实现需求

原始数据 #

import pandas as pd
import numpy as np

# 示例输入DataFrame
df = pd.DataFrame({
    'Fruits': ['Apple', 'Orange', 'Grape', 'Berry', 'Apple'],
    'PRS_001_Person_ABCD': [0.5, 0.2, np.nan, np.nan, np.nan],
    'PRS_002_Person_ABCD': [1.3, np.nan, 0.06, np.nan, 1.3],
    'PRS_015_Person_ABCD': [np.nan, 0.021, np.nan, 0.3, 0.5],
    'PRS_016_Person_ABCD': [np.nan, np.nan, 0.7, 0.04, np.nan]
})

目标结构 #

# 期望输出的DataFrame结构
target_df = pd.DataFrame({
    'Fruits': ['Apple', 'Orange', 'Grape', 'Berry', 'Apple'],
    'Value1': [0.5, 0.2, 0.06, 0.3, 1.3],
    'Value2': [1.3, 0.021, 0.7, 0.04, 0.5],
    'Person1': ['Product 1', 'Product 1', 'Product 2', 'Product 15', 'Product 2'],
    'Person2': ['Product 2', 'Product 15', 'Product 16', 'Product 16', 'Product 15']
})

实现代码 #

步骤1：转长表并清理缺失值 #

将宽格式的列转换为行，同时过滤掉无有效值的行：

# 转长表，保留Fruits作为标识列
melted = df.melt(id_vars='Fruits', var_name='Person_Code', value_name='Value')
# 移除缺失值行（每组Fruits将保留2行有效数据）
melted = melted.dropna(subset=['Value'])

步骤2：生成Person名称 #

从Person_Code中提取产品编号，格式化为Product X：

# 提取PRS后的数字部分，去除前导零后生成Person名称
melted['Person'] = melted['Person_Code'].str.extract(r'PRS_(\d+)_').apply(
    lambda x: f"Product {int(x[0])}", axis=1
)

步骤3：重组为目标结构 #

按Fruits分组，将每组的两个有效条目分别映射到Value1/Value2和Person1/Person2：

# 分组聚合生成目标列
result = melted.groupby('Fruits', as_index=False).agg(
    Value1=('Value', lambda x: x.iloc[0]),
    Value2=('Value', lambda x: x.iloc[1]),
    Person1=('Person', lambda x: x.iloc[0]),
    Person2=('Person', lambda x: x.iloc[1])
)

# 调整列顺序与目标一致
result = result[['Fruits', 'Value1', 'Value2', 'Person1', 'Person2']]

验证结果 #

运行后result将完全匹配目标结构，可直接用于后续分析。

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