Python如何生成满足指定约束、元素和为1的固定长度随机数组

实现思路 #

核心逻辑是先固定满足约束的元素取值，再分配剩余额度，全程不会出现不符合要求的中间结果，无需筛选：

• 第一步：先处理有上限约束的0标记元素，直接在各自上限范围内随机生成取值，从生成之初就满足上限要求，后续不再修改。
• 第二步：给有下限约束的2标记元素预留下限值，确保后续无论加多少增量，最终取值都大于要求的下限。
• 第三步：计算总固定占用额度（下限值+两个0标记元素的取值），如果该额度已经大于等于1，说明给定的约束参数不合理，直接抛出异常。
• 第四步：将剩余可分配额度，随机拆分为1（2标记元素的增量） + 1标记元素数量份正数，分别分配给对应位置，最终所有元素之和刚好为1。

完整代码实现 #

依赖Numpy版本（推荐，生成的随机数分布更均匀） #

import numpy as np

def generate_constrained_array(higher_than: float, lower_than_1: float, lower_than_2: float, mask: list) -> list:
    # 定位各类型元素的索引位置
    pos_2 = mask.index(2)
    pos_0_list = [i for i, val in enumerate(mask) if val == 0]
    pos_1_list = [i for i, val in enumerate(mask) if val == 1]

    res = [0.0] * len(mask)
    # 生成两个0标记位置的取值，直接满足上限约束
    res[pos_0_list[0]] = np.random.uniform(0, lower_than_1)
    res[pos_0_list[1]] = np.random.uniform(0, lower_than_2)

    # 计算已占用额度和剩余可分配额度
    used = higher_than + res[pos_0_list[0]] + res[pos_0_list[1]]
    if used >= 1:
        raise ValueError("约束参数不合理，固定占用额度已超过总和1，请调整约束值")
    remaining = 1 - used

    # 拆分剩余额度：1份给2标记的增量，其余给1标记元素
    n_split = len(pos_1_list) + 1
    # Dirichlet分布生成和为1的随机数，乘以剩余额度得到各份大小
    split_vals = np.random.dirichlet(np.ones(n_split)) * remaining

    # 赋值2标记和1标记元素
    res[pos_2] = higher_than + split_vals[0]
    for idx, pos in enumerate(pos_1_list):
        res[pos] = split_vals[idx + 1]
    
    return res

# 测试调用
if __name__ == "__main__":
    higher_than = 0.4
    lower_than_1 = 0.2
    lower_than_2 = 0.1
    mask = [1, 2, 0, 1, 1, 1, 0]
    arr = generate_constrained_array(higher_than, lower_than_1, lower_than_2, mask)
    print("生成数组：", [round(i, 4) for i in arr])
    print("数组和：", round(sum(arr), 4))
    print("第二个元素是否大于0.4：", arr[1] > 0.4)
    print("第三个元素是否小于0.2：", arr[2] < 0.2)
    print("第七个元素是否小于0.1：", arr[6] < 0.1)

无第三方依赖版本 #

如果不想引入Numpy，可以用均匀切割的方式实现剩余额度拆分：

import random

def generate_constrained_array_no_numpy(higher_than: float, lower_than_1: float, lower_than_2: float, mask: list) -> list:
    pos_2 = mask.index(2)
    pos_0_list = [i for i, val in enumerate(mask) if val == 0]
    pos_1_list = [i for i, val in enumerate(mask) if val == 1]

    res = [0.0] * len(mask)
    res[pos_0_list[0]] = random.uniform(0, lower_than_1)
    res[pos_0_list[1]] = random.uniform(0, lower_than_2)

    used = higher_than + res[pos_0_list[0]] + res[pos_0_list[1]]
    if used >= 1:
        raise ValueError("约束参数不合理，固定占用额度已超过总和1，请调整约束值")
    remaining = 1 - used

    n_split = len(pos_1_list) + 1
    # 均匀切割实现随机拆分
    cuts = sorted([random.uniform(0, remaining) for _ in range(n_split - 1)])
    split_vals = [cuts[0]] + [cuts[i] - cuts[i-1] for i in range(1, n_split-1)] + [remaining - cuts[-1]]

    res[pos_2] = higher_than + split_vals[0]
    for idx, pos in enumerate(pos_1_list):
        res[pos] = split_vals[idx + 1]
    
    return res

验证说明 #

每次调用函数都会直接生成符合所有要求的数组：

1. 长度和输入的mask数组一致，示例中为7
2. 2标记元素取值=预留下限+正增量，必然大于higher_than
3. 两个0标记元素直接在上限范围内生成，必然满足小于对应上限的要求
4. 所有元素加和刚好为1，无精度外的误差

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