将groupby操作生成的Series赋值给DataFrame列时出现NaN的问题求助

问题场景 #

数据定义：

import pandas as pd

df_SOT = pd.DataFrame({
    'Lane': {26055: 'L2', 26056: 'L2', 26057: 'L2', 26058: 'L2', 26059: 'L2', 
             25972: 'L1', 25973: 'L1', 25974: 'L1', 25975: 'L1', 25976: 'L1'},
    'Carrier SCAC': {26055: 'JNJR', 26056: 'WOSQ', 26057: 'BGME', 26058: 'ITSB', 26059: 'UCSB',
                    25972: 'BGME', 25973: 'SCNN', 25974: 'XPOL', 25975: 'SJRG', 25976: 'MTRK'},
    'Annual Volume': {26055: 5604.0, 26056: 5604.0, 26057: 5604.0, 26058: 5604.0, 26059: 5604.0,
                     25972: 4917.0, 25973: 4917.0, 25974: 4917.0, 25975: 4917.0, 25976: 4917.0},
    'Annual Capacity': {26055: 260.0, 26056: 1300.0, 26057: 2704.0, 26058: 2080.0, 26059: 4368.0,
                       25972: 5408.0, 25973: 3380.0, 25974: 4940.0, 25975: 156.0, 25976: 4940.0}
})

定义分配函数：

def allocation(df_alloc):
    Annual_Volume = df_alloc['Annual Volume']
    Annual_Capacity = df_alloc['Annual Capacity']
    Allocation = []
    Cum_Capacity = 0
    for idx in df_alloc.index:
        Allocate = min(0.5*Annual_Volume[idx], Annual_Capacity[idx], Annual_Volume[idx]-Cum_Capacity)
        Cum_Capacity += Allocate
        Allocation.append(Allocate)
    return Allocation

执行groupby操作后得到结果：

df_SOT.groupby('Lane').apply(allocation)

输出：

Lane
L1         [2458.5, 2458.5, 0.0, 0.0, 0.0]
L2    [260.0, 1300.0, 2704.0, 1340.0, 0.0]
dtype: object

尝试赋值到原DataFrame时出现问题：

df_SOT['Allocation'] = df_SOT.groupby('Lane').apply(allocation)

此时Allocation列全为NaN。

问题原因 #

groupby('Lane').apply(allocation)返回的Series索引是分组键['L1','L2']，但原DataFrame的索引是26055、26056等行号，两者索引无法对齐，导致赋值时所有位置都匹配不到对应值，最终填充为NaN。

解决方案 #

方案1：修改函数返回带原组索引的Series #

调整allocation函数，让它返回带有原分组行索引的Series，这样apply后会自动对齐原DataFrame的索引：

def allocation(df_alloc):
    Annual_Volume = df_alloc['Annual Volume']
    Annual_Capacity = df_alloc['Annual Capacity']
    Allocation = []
    Cum_Capacity = 0
    for idx in df_alloc.index:
        Allocate = min(0.5*Annual_Volume[idx], Annual_Capacity[idx], Annual_Volume[idx]-Cum_Capacity)
        Cum_Capacity += Allocate
        Allocation.append(Allocate)
    # 返回带原分组行索引的Series
    return pd.Series(Allocation, index=df_alloc.index)

重新执行赋值：

df_SOT['Allocation'] = df_SOT.groupby('Lane').apply(allocation)

此时Allocation列会正确填充对应数值。

方案2：向量化实现（更优） #

避免循环，用Pandas的向量化操作提升效率，尤其适合大数据量场景：

def allocation_vectorized(df_alloc):
    half_volume = 0.5 * df_alloc['Annual Volume'].iloc[0]
    # 计算每个位置的剩余可分配量：总半量 - 前面的累积分配
    cum_capacity = df_alloc['Annual Capacity'].cumsum().shift(1, fill_value=0)
    remaining = half_volume - cum_capacity
    # 取三个值的最小值：年度容量、剩余可分配量、半量
    allocation = pd.DataFrame({
        'capacity': df_alloc['Annual Capacity'],
        'remaining': remaining
    }).min(axis=1)
    # 确保累积分配不超过总半量
    cum_allocation = allocation.cumsum()
    excess = cum_allocation - half_volume
    allocation = allocation.where(excess <= 0, allocation - excess)
    return allocation

df_SOT['Allocation'] = df_SOT.groupby('Lane').apply(allocation_vectorized)

这个实现逻辑和原函数一致，但去掉了循环，利用Pandas的内置方法提升计算效率。

内容的提问来源于stack exchange，提问作者Rajib Lochan Sarkar