Python为DataFrame生成满足总和与单值上限双约束的随机整数

前置校验 #

首先需要确认存在可行解：所有行CONSTRAINT列的总和必须大于等于目标总和（本案例为300），且目标总和不能为负数，否则不存在符合要求的结果。

实现逻辑 #

采用逐次随机分配的逻辑，既保证随机性，又100%符合所有约束：

1. 初始化新增列RANDOM所有值为0，剩余待分配总额为目标值300
2. 循环分配直到待分配总额归零：
   • 筛选当前取值未达到CONSTRAINT上限的可分配行
   • 随机选中一行，为其分配1到「该行剩余可分配额度、剩余待分配总额二者最小值」之间的随机整数
   • 扣除对应待分配额度
3. 最终生成的列自动满足总和固定、不超上限、非负整数的所有要求

完整代码示例 #

import pandas as pd
import random

# 替换为你自己的DataFrame即可
df = pd.DataFrame({
    'GEOID': ['010010000001', '010010000002', '010010000003', '010010000004', '010010000005'],
    'CONSTRAINT': [100, 50, 75, 75, 100]
})

target_total = 300

# 可行解校验
if df['CONSTRAINT'].sum() < target_total or target_total < 0:
    raise ValueError("约束总和小于目标总和或目标总和为负，无可行解")

# 初始化随机列
df['RANDOM'] = 0
remaining_total = target_total

while remaining_total > 0:
    # 获取可分配的行索引
    eligible_idx = df[df['RANDOM'] < df['CONSTRAINT']].index.tolist()
    # 随机选行
    selected_idx = random.choice(eligible_idx)
    # 计算本次最大可分配数值
    max_add = min(
        df.loc[selected_idx, 'CONSTRAINT'] - df.loc[selected_idx, 'RANDOM'],
        remaining_total
    )
    # 随机生成本次分配数值，如需提升运行效率可直接赋值max_add大幅减少循环次数
    add_val = random.randint(1, max_add)
    df.loc[selected_idx, 'RANDOM'] += add_val
    remaining_total -= add_val

结果验证 #

运行后可以自行校验约束是否全部满足：

• print(df['RANDOM'].sum()) 确认总和等于300
• print((df['RANDOM'] <= df['CONSTRAINT']).all()) 确认所有值不超过对应上限
• print((df['RANDOM'] >= 0).all()) 确认所有值非负

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