如何用Pandas向量化函数替代循环实现数值按规则分配？

数据背景 #

现有DataFrame：

import pandas as pd
df = pd.DataFrame(data={'Name': ['a1', 'a2','a1'], 'Quantity': [8,4,5], 'Order': [1,1,2], 'Result': [0,0,0]})

对应表格：

分配规则 #

需将数值N=10按以下规则分配至Result列：

• 每行分配值不超过该行的Quantity；
• 按Order顺序优先满足同Name分组内的行。

期望结果 #

现有解法 #

1. 按Name分组获取分组键；
2. 嵌套循环遍历Order和Name，直至N耗尽。

向量化解决方案 #

完全可以用Pandas向量化操作替代循环，核心思路是先按规则排序确定分配优先级，再通过累计量计算每行的实际分配值，全程无需循环：

步骤1：按优先级排序 #

先按Order升序、Name排序，确保分配顺序符合规则：

df_sorted = df.sort_values(by=['Order', 'Name'], ignore_index=True)

步骤2：计算累计可分配上限 #

计算每行及之前所有行的累计Quantity，用于判断当前行的分配额度：

df_sorted['cum_quantity'] = df_sorted['Quantity'].cumsum()
prev_cum = df_sorted['cum_quantity'] - df_sorted['Quantity']  # 前一行的累计量

步骤3：向量化计算分配值 #

通过掩码区分三种分配情况，直接批量赋值：

N = 10

# 三种情况的掩码
mask_full = df_sorted['cum_quantity'] <= N  # 累计量未超N，全部分配
mask_partial = (prev_cum < N) & (~mask_full)  # 前序累计未超N，当前累计超N，分配剩余额度
mask_zero = prev_cum >= N  # 前序累计已超N，无分配

# 批量赋值
df_sorted.loc[mask_full, 'Result'] = df_sorted['Quantity']
df_sorted.loc[mask_partial, 'Result'] = N - prev_cum
df_sorted.loc[mask_zero, 'Result'] = 0

步骤4：恢复原行顺序 #

如果需要保持原始DataFrame的行顺序，通过合并还原：

df = df.merge(
    df_sorted[['Name', 'Order', 'Quantity', 'Result']],
    on=['Name', 'Order', 'Quantity'],
    suffixes=('', '_new')
).drop(columns=['Result_new'])

完整代码示例 #

import pandas as pd

df = pd.DataFrame(data={'Name': ['a1', 'a2','a1'], 'Quantity': [8,4,5], 'Order': [1,1,2], 'Result': [0,0,0]})
N = 10

# 排序确定分配优先级
df_sorted = df.sort_values(by=['Order', 'Name'], ignore_index=True)
# 计算累计量
df_sorted['cum_quantity'] = df_sorted['Quantity'].cumsum()
prev_cum = df_sorted['cum_quantity'] - df_sorted['Quantity']

# 向量化计算Result
mask_full = df_sorted['cum_quantity'] <= N
mask_partial = (prev_cum < N) & (~mask_full)
mask_zero = prev_cum >= N

df_sorted.loc[mask_full, 'Result'] = df_sorted['Quantity']
df_sorted.loc[mask_partial, 'Result'] = N - prev_cum
df_sorted.loc[mask_zero, 'Result'] = 0

# 恢复原顺序
df = df.merge(
    df_sorted[['Name', 'Order', 'Quantity', 'Result']],
    on=['Name', 'Order', 'Quantity'],
    suffixes=('', '_new')
).drop(columns=['Result_new'])

print(df)

运行输出结果与期望一致：

Name  Quantity  Order  Result
0   a1         8      1       8
1   a2         4      1       2
2   a1         5      2       0

内容的提问来源于stack exchange，提问作者Caterina De Franco