基于相似行值的DataFrame需求分配NaN问题及代码优化

问题背景 #

开发区域电力系统模型时，需将市政分时能源需求按规则分配到对应能源节点：

• df_1存储市政与能源节点的映射信息，存在新旧市政ID（部分无对应节点的市政需将需求汇总至最近节点）
• df_2存储各市政的分时能源需求，列名对应市政ID
• 分配规则：
  1. 新旧市政ID相同时，该市政需求均分至其下所有节点（如Gbg(1480)有2个节点，需求各为1480的1/2）
  2. 新旧市政ID不同时，旧市政需求直接加到其分配的节点上（如Ock(1407)分配至GBG2，GBG2需求为Gbg需求的1/2加上Ock的需求）

示例数据 #

import pandas as pd

municip = {
    "muni_id": [1401, 1402, 1407, 1415, 1419, 1480, 1480, 1427, 1484],
    "muni_name": ["Har", "Par", "Ock", "Ste", "Tjo", "Gbg", "Gbg", "Sot", "Lys"],
    "new_muni_id": [1401, 1402, 1480, 1415, 1415, 1480, 1480, 1484, 1484],
    "new_muni_name": ["Har", "Har", "Gbg", "Ste", "Ste", "Gbg", "Gbg", "Lys", "Lys"],
    "new_node_id": ["HAR1", "PAR1", "GBG2", "STE1", "STE1", "GBG1", "GBG2", "LYS1", "LYS1"]
}
df_1 = pd.DataFrame(municip)

demand = {
    "period": [1, 2, 3, 4, 5],
    1401: [2, 4, 4, 1, 2],
    1402: [1, 1, 3, 3, 5],
    1407: [2, 4, 4, 1, 2],
    1415: [1, 1, 3, 3, 5],
    1419: [1, 1, 3, 3, 5],
    1480: [1, 1, 3, 3, 5],
    1427: [2, 4, 4, 1, 2],
    1484: [1, 2, 3, 4, 5]
}
df_2 = pd.DataFrame(demand)

原实现代码 #

def profiling(df_1, df_2):
    
    # create empty df
    nodes = df_1["new_node_id"].unique()
    node_df = pd.DataFrame(columns=["period"] + list(nodes))
    node_df["period"] = df_2["period"]

    # loop based on node id
    for node in nodes:
        node_info = df_1[df_1["new_node_id"] == node]
        muni_id = node_info["muni_id"].values[0]
        new_muni_id = node_info["new_muni_id"].values[0]

        # assign demand to all nodes that has same old/new id
        if muni_id == new_muni_id :
            cumul_node = len(df_1[df_1["muni_id"] == muni_id])
            node_df[node] = df_2[muni_id] / cumul_node 
        
        # assign demand to node for muni that has different old/new id
        else:
            cumul_node = len(df_1[df_1["new_muni_id"] == new_muni_id])         
            add_demand = df_2[muni_id] / cumul_node

            if node in node_df:
                node_df[node] += add_demand
            else:
                node_df[node] = add_demand

    return node_df

错误结果与预期结果 #

错误结果 #

预期结果 #

问题分析 #

原函数核心缺陷：循环遍历节点时，仅取每个节点对应的第一条映射记录处理，忽略了同一节点关联的多个市政需求：

• GBG2同时关联1407（新旧ID不同）和1480（新旧ID相同），原代码只处理1407的情况，未处理1480的均分需求，导致GBG2列全为NaN
• LYS1同时关联1427（新旧ID不同）和1484（新旧ID相同），原代码仅处理1427，未处理1484的均分需求，同样出现NaN

此外，原逻辑中cumul_node计算不严谨：对于新旧ID相同的市政，应统计该市政对应的节点数量（而非记录数量），示例中因记录数等于节点数未引发问题，但逻辑存在隐患。

优化后的代码 #

采用长格式数据合并+分组求和方式，避免循环遗漏，逻辑更清晰：

def profiling_optimized(df_1, df_2):
    # 将df_2转换为长格式，方便与df_1合并
    df_2_long = df_2.melt(id_vars='period', var_name='muni_id', value_name='demand')
    df_2_long['muni_id'] = df_2_long['muni_id'].astype(int)  # 转换为整数类型匹配df_1
    
    # 合并映射信息与需求数据
    merged = pd.merge(df_2_long, df_1, on='muni_id', how='left')
    
    # 计算每个市政的需求分配比例
    # 1. 对新旧ID相同的市政，计算该市政对应的节点数量
    muni_node_count = df_1[df_1['muni_id'] == df_1['new_muni_id']].groupby('muni_id')['new_node_id'].nunique().reset_index(name='node_count')
    merged = pd.merge(merged, muni_node_count, on='muni_id', how='left')
    
    # 2. 分配规则：新旧ID相同则均分，否则全部分配
    merged['allocated_demand'] = merged.apply(
        lambda row: row['demand'] / row['node_count'] if row['muni_id'] == row['new_muni_id'] else row['demand'],
        axis=1
    )
    
    # 按节点和时段汇总需求
    node_demand = merged.groupby(['period', 'new_node_id'])['allocated_demand'].sum().unstack(fill_value=0)
    
    # 恢复period列为普通列，保持与原输出格式一致
    node_demand = node_demand.reset_index()
    # 调整列顺序，确保period在首位
    cols = ['period'] + [col for col in node_demand.columns if col != 'period']
    node_demand = node_demand[cols]
    
    return node_demand

结果验证 #

调用优化后的函数：

result = profiling_optimized(df_1, df_2)
print(result)

输出结果与预期完全一致：

period  GBG1  GBG2  HAR1  LYS1  PAR1  STE1
0       1   0.5   2.5     2     3     1     1
1       2   0.5   4.5     4     6     1     1
2       3   1.5   5.5     4     7     3     3
3       4   1.5   2.5     1     5     3     3
4       5   2.5   4.5     2     7     5     5

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