使用Apply与Iterrows为DataFrame按组分配Weight列遇错求助

输入数据 #

import pandas as pd

data=pd.DataFrame({'Name':['a','b','c','d','e','f','g','h','i'],
              'Group':['G1','G3','G2','G1','G2','G3','G1','G2','G3'],
              'Value':[15,21,14,8,17,22,11,13,15]})

data=data.sort_values('Value',ascending=False)

需求说明 #

• 先对Value列降序排序；
• 为每行创建Weight列，值为1或0：每个Group对应一个预算阈值，按排序后的顺序累计Value，若累计值加上当前行Value超过阈值，当前行及该组剩余行的Weight设为0，否则设为1并累计Value。

预期输出 #

data_out=pd.DataFrame({'Name':['a','b','c','d','e','f','g','h','i'],
              'Group':['G1','G3','G2','G1','G2','G3','G1','G2','G3'],
              'Value':[15,21,14,8,17,22,11,13,15],
              'Weight':[1,0,1,0,1,1,1,0,0]})

实现逻辑示例 #

以G1为例，阈值为27：

• a的Value为15，累计15≤27，Weight=1；
• g的Value为11，累计26≤27，Weight=1；
• d的Value为8，累计会超27，Weight=0。
  G2、G3同理。

尝试的代码及问题 #

以下使用apply的代码报错：

threshold_dic={'G1':27,'G2':32,'G3':25}
initial_dic={'G1':0,'G2':0,'G3':0}
def f(row):
    if initial_dic[row['Group']]<= threshold_dic[row['Group']]:
        row['Weight']=1
    else:
        row['Weight']=0
    initial_dic[row['Group']]+=row['Value']
    return row
data.apply(f,axis=0)

改用iterrows循环仍无法正常运行。

解决方案 #

问题核心原因：

1. apply使用axis=0是按列处理，应改为axis=1按行处理，但更关键的是，依赖全局变量的方式在apply中可能因pandas内部优化导致处理顺序混乱，累计值计算错误；
2. 直接修改排序后的数据时，未考虑恢复原数据的行顺序，导致最终结果不符合预期。

正确实现代码：

import pandas as pd

# 输入数据
data=pd.DataFrame({'Name':['a','b','c','d','e','f','g','h','i'],
              'Group':['G1','G3','G2','G1','G2','G3','G1','G2','G3'],
              'Value':[15,21,14,8,17,22,11,13,15]})

# 按Value降序排序，保留原索引用于后续恢复顺序
data_sorted = data.sort_values('Value', ascending=False).copy()
threshold_dic = {'G1':27, 'G2':32, 'G3':25}
# 初始化每组的累计值
current_sum = {group: 0 for group in threshold_dic.keys()}
weights = []

# 遍历排序后的行，逐个计算Weight
for _, row in data_sorted.iterrows():
    group = row['Group']
    val = row['Value']
    if current_sum[group] + val <= threshold_dic[group]:
        weights.append(1)
        current_sum[group] += val
    else:
        weights.append(0)

# 赋值Weight列并恢复原数据的行顺序
data_sorted['Weight'] = weights
data_out = data_sorted.loc[data.index].reset_index(drop=True)

print(data_out)

运行后即可得到与预期一致的结果。

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