pandas - 生成具有条件的随机数列表 - numpy

Question

这个问题已经有答案了:

Is there an efficient way to generate N random integers in a range that have a given sum or average? (6 个回答)

已关闭 3 年前。

我想生成一个包含 15 个整数的列表，总和为 12，最小值为 0，最大值为 6。

我尝试了以下代码

def generate(low,high,total,entity):
   while sum(entity)!=total:
       entity=np.random.randint(low, high, size=15)
   return entity

但是上述功能无法正常工作。太费时间了。请告诉我生成此类数字的有效方法？

Answer 1

严格来说，上述内容是有效的。但对于 0 到 6 之间的 15 个数字，生成 12 的几率并不高。事实上，我们可以通过以下方式计算可能性的数量:

F(s, 1) = 1 对于0≤s≤6和

F(s, n) = Σ⁶_i=0F(s-i, n-1)。

我们可以用一个值来计算:

from functools import lru_cache

@lru_cache()
def f(s, n, mn, mx):
    if n < 1:
        return 0
    if n == 1:
        return int(mn <= s <= mx)
    else:
        if s < mn:
            return 0
        return sum(f(s-i, n-1, mn, mx) for i in range(mn, mx+1))

这意味着在 4'747'561'509'943 种可能性中，有 9'483'280 种可能性产生总和为 12，即 0.00019975%。因此，需要大约 500'624 次迭代才能得出这样的解决方案。

因此，我们应该更好地致力于找到一种直接的方法来生成这样的序列。我们可以通过每次计算生成数字的概率来做到这一点:生成 i 作为 n 个数字序列中第一个数字的概率，总和为 < em>s 为 F(s-i, n-1, 0, 6)/F(s, n, 0, 6)。这将保证我们在可能性列表上生成一个均匀列表，如果我们每次都绘制一个均匀的数字，那么它将不匹配与给定值匹配的整个值列表上的均匀分布条件:

我们可以递归地做到这一点:

from numpy import choice

def sumseq(n, s, mn, mx):
    if n > 1:
        den = f(s, n, mn, mx)
        val, = choice(
            range(mn, mx+1),
            1,
            p=[f(s-i, n-1, mn, mx)/den for i in range(mn, mx+1)]
        )
        yield val
        yield from sumseq(n-1, s-val, mn, mx)
    elif n > 0:
        yield s

通过上面的函数，我们可以生成numpy数组:

>>> np.array(list(sumseq(15, 12, 0, 6)))
array([0, 0, 0, 0, 0, 4, 0, 3, 0, 1, 0, 0, 1, 2, 1])
>>> np.array(list(sumseq(15, 12, 0, 6)))
array([0, 0, 1, 0, 0, 1, 4, 1, 0, 0, 2, 1, 0, 0, 2])
>>> np.array(list(sumseq(15, 12, 0, 6)))
array([0, 1, 0, 0, 2, 0, 3, 1, 3, 0, 1, 0, 0, 0, 1])
>>> np.array(list(sumseq(15, 12, 0, 6)))
array([5, 1, 0, 0, 0, 0, 1, 1, 0, 0, 1, 1, 1, 0, 1])
>>> np.array(list(sumseq(15, 12, 0, 6)))
array([0, 0, 0, 0, 4, 2, 3, 0, 0, 0, 0, 0, 3, 0, 0])