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

Question

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

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

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

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

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()) 函数只接受整数作为参数。}

Answer 1

这是一个精确的（每个合法总和都有相同的概率）解决方案。它使用所有合法总和的枚举，不是在我们遍历每个总和的意义上，而是给定一个数字 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。

Answer 2

修改了你的代码以支持数组，

>>> def foo(a, b):
    p = np.ones(b) / b
    arr = []
    if isinstance(a, type(np.array([]))):
        for i in a:
            arr.extend(np.random.multinomial(i, p, size = 1))
        return np.array(arr)
    return np.random.multinomial(a, p, size = 1)

---

输出：

>>> foo(np.array([12, 18, 21]), 5)
array([[3, 4, 3, 1, 1],
       [5, 4, 3, 4, 2],
       [3, 5, 3, 6, 4]])

Answer 3

您可以使用多项式进行如下操作，其中np.random.multinomial(s, [1/5]*5) 类似于绘制一个 5 面 s 次的无偏骰子，它会返回每面出现在每次抽签中的次数。所以每次出现的计数总和必须等于总抽奖次数

def foo(arr, n):
    return np.r_[[np.random.multinomial(s, [1/5]*5) for s in arr]]

foo(np.array([12, 18, 21]), 5)

Answer 4

以下与我的其他答案的不同之处在于它保证了所有可能输出的均匀分布，但可能仅适用于小输入（可能输出的空间相当小）：

import random

def gen_all(low, high, n, total):
  """Yields all possible n-tuples of [low; high] ints that sum up to total."""
  if n == 0:
    yield ()
    return
  adj_low = max(low, total - (n - 1) * high)
  adj_high = min(high, total - (n - 1) * low)
  for val in range(adj_low, adj_high + 1):
    for arr in gen_all(low, high, n - 1, total - val):
      yield (val,) + arr

def gen(low, high, n, total):
  return random.choice(list(gen_all(low, high, n, total)))

print(gen(low=1, high=5, n=5, total=5))
print(gen(low=1, high=5, n=5, total=12))
print(gen(low=1, high=5, n=5, total=18))
print(gen(low=1, high=5, n=5, total=25))

这会将所有可能的输出具体化到一个列表中，然后再选择一个。在实践中，可能会使用reservoir sampling（大小为 1）。

如果您需要生成多个具有相同参数的随机集，可以使用带替换的水库采样：https://epubs.siam.org/doi/pdf/10.1137/1.9781611972740.53

Answer 5

这是一个可能比其他解决方案更脏的解决方案。但是只要random.sample给它，它就会给它均匀分布。

from random import sample

def foo(num, count):
    result = [1] * count
    indices = set(range(count))

    for _ in range(num - count):
        idx = sample(indices, 1)[0]
        result[idx] += 1

        if result[idx] == 5:
            indices.remove(idx)

    return result

print(foo(16, 4))
# [4, 4, 5, 3]

它以ones 列表开始，不断添加+1 直到达到目标，还跟踪哪个值达到5，因此不再添加。非常基本且缓慢的解决方案。 （还是O(n)）

编辑：我已经为单个值编写了它，您必须在循环中应用它以获得多个结果。

Answer 6

这是一个快速而肮脏的版本（不容易矢量化，并且不保证可能的输出有任何特定分布）：

import random

def gen(low, high, num, total):
  arr = []
  while num > 0:
    adj_low = max(low, total - (num - 1) * high)
    adj_high = min(high, total - (num - 1) * low)
    val = random.randint(adj_low, adj_high)
    arr.append(val)
    total -= val
    num -= 1
  random.shuffle(arr)
  return arr

print(gen(low=1, high=5, num=5, total=5))
print(gen(low=1, high=5, num=5, total=12))
print(gen(low=1, high=5, num=5, total=18))
print(gen(low=1, high=5, num=5, total=25))

它一次选择一个随机整数。

在每一点上，它都知道有多少total 需要（随机）分配给剩余的数字（num）。它使用这些知识来调整 low 和 high 边界以始终保持在约束范围内。

此方法使稍后选择的数字偏向于具有更大的值。 shuffle 消除了这种偏差（但同样，它不提供任何关于可能输出的整体分布的保证）。

我很高兴听到是否有更好的方法来做到这一点。