时间戳

首页 > 解决方案 > 生成已知总数的随机整数数组

python - 生成已知总数的随机整数数组

问题描述

考虑以下玩具数组，其整数范围为5- 25：

a = np.array([12, 18, 21])

如何生成5范围从1-的随机整数数组5，其总和等于数组中的每个数字a？该解决方案应在所有可能的输出上生成均匀分布。

到目前为止，我已经设法创建了一个简单的函数，它将产生5随机整数：

import numpy as np

def foo(a, b):
    p = np.ones(b) / b
    return np.random.multinomial(a, p, size = 1)

使用数组a* 中的值的示例：

In [1]: foo(12, 5)
Out[1]: array([[1, 4, 3, 2, 2]])

In [2]: foo(18, 5)
Out[2]: array([[2, 2, 3, 3, 8]])

In [3]: foo(21, 5)
Out[3]: array([[6, 5, 3, 4, 3]])

显然，这些整数具有所需的总数，但它们并不总是介于1和之间5。

预期的输出将是以下几行：

In [4]: foo(np.array([12, 18, 21]), 5)
Out[4]: 
array([[1, 4, 3, 2, 2],
       [4, 3, 3, 3, 5],
       [5, 5, 4, 4, 3]])

* np.multinomial()) 函数只接受整数作为参数。

标签： pythonpython-3.xnumpy

解决方案

这是一个精确的（每个合法金额都有相同的概率）解决方案。它使用所有合法总和的枚举，不是在我们遍历每个总和的意义上，而是给定一个数字 n，我们可以直接计算枚举中的第 n 个总和。由于我们也知道合法和的总数，我们可以简单地绘制统一整数并对其进行转换：

import numpy as np
import functools as ft

#partition count
@ft.lru_cache(None)
def capped_pc(N,k,m):
    if N < 0:
        return 0
    if k == 0:
        return int(N<=m)
    return sum(capped_pc(N-i,k-1,m) for i in range(m+1))

capped_pc_v = np.vectorize(capped_pc)

def random_capped_partition(low,high,n,total,size=1):
    total = total - n*low
    high = high - low
    if total > n*high or total < 0:
        raise ValueError
    idx = np.random.randint(0,capped_pc(total,n-1,high),size)
    total = np.broadcast_to(total,(size,1))
    out = np.empty((size,n),int)
    for j in range(n-1):
        freqs = capped_pc_v(total-np.arange(high+1),n-2-j,high)
        freqs_ps = np.c_[np.zeros(size,int),freqs.cumsum(axis=1)]
        out[:,j] = [f.searchsorted(i,"right") for f,i in zip(freqs_ps[:,1:],idx)]
        idx = idx - np.take_along_axis(freqs_ps,out[:,j,None],1).ravel()
        total = total - out[:,j,None]
    out[:,-1] = total.ravel()
    return out + low

演示：

# 4 values between 1 and 5 summing to 12
# a million samples takes a few seconds
x = random_capped_partition(1,5,4,12,1000000)

# sanity checks

# values between 1 and 5
x.min(),x.max()
# (1, 5)

# all legal sums occur
# count them brute force
sum(1 for a in range(1,6) for b in range(1,6) for c in range(1,6) if 7 <=     a+b+c <= 11)
# 85
# and count unique samples
len(np.unique(x,axis=0))
# 85

# check for uniformity
np.unique(x, axis=0, return_counts=True)[1]
# array([11884, 11858, 11659, 11544, 11776, 11625, 11813, 11784, 11733,
#        11699, 11820, 11802, 11844, 11807, 11928, 11641, 11701, 12084,
#        11691, 11793, 11857, 11608, 11895, 11839, 11708, 11785, 11764,
#        11736, 11670, 11804, 11821, 11818, 11798, 11587, 11837, 11759,
#        11707, 11759, 11761, 11755, 11663, 11747, 11729, 11758, 11699,
#        11672, 11630, 11789, 11646, 11850, 11670, 11607, 11860, 11772,
#        11716, 11995, 11802, 11865, 11855, 11622, 11679, 11757, 11831,
#        11737, 11629, 11714, 11874, 11793, 11907, 11887, 11568, 11741,
#        11932, 11639, 11716, 12070, 11746, 11787, 11672, 11643, 11798,
#        11709, 11866, 11851, 11753])

简要说明：

我们使用简单的递归来计算上限分区的总数。我们在第一个 bin 上进行拆分，即我们固定第一个 bin 中的数字，并通过递归检索填充剩余 bin 的方法数。然后我们简单地总结不同的第一个 bin 选项。我们使用缓存装饰器来控制递归。这个装饰器会记住我们已经完成的所有参数组合，因此当我们通过不同的递归路径到达相同的参数时，我们不必再次进行相同的计算。

枚举的工作方式类似。假设字典顺序。如何找到第 n 个分区？再次，在第一个垃圾箱上拆分。因为我们知道对于每个值，第一个 bin 可以采用剩余 bin 可以填充的方式的总数，我们可以形成累积和，然后查看 n 适合的位置。如果 n 位于第 i 个和第 i+1 个部分和之间第一个索引是 i+low，我们从 n 中减去第 i 个和并从剩余的 bin 开始。

推荐阅读