在 k 个容器中随机分配一个整数，并在分配中均匀分布

Question

我正在寻找一种高效的 Python 函数，它可以在 k 个容器中随机分配一个整数。也就是说，某些函数 allocate(n, k) 将生成一个 k 大小的整数数组，总和为 n.

例如，allocate(3, 2) 会以相同的概率产生 [3, 0]、[2, 1]、[1, 2] 或 [0, 3]（不同于，分配，而不是项目，应该均匀分布）。

Answer 1

使用“星条形”方法，我们可以将其转化为从 n+k-1 个可能位置的列表中为可能的分隔符选择 k-1 个位置的问题。 (Wikipedia proof)

from random import sample

def allocate(n,k):
    dividers = sample(range(1, n+k), k-1)
    dividers = sorted(dividers)
    dividers.insert(0, 0)
    dividers.append(n+k)
    return [dividers[i+1]-dividers[i]-1 for i in range(k)]
    
print(allocate(4,3))

有 ((n+k-1) choose (k-1)) 种可能的分配符合您的标准，每一种分配都对应于选择 k 个位置的特定方式，将分隔符放在 n+k-1 个位置中对于对象，这同样可能导致每个对象。

（请注意与 proposed-in-comments existing answer to a similar question 的细微差别：此问题要求 non-negative 整数的有序序列，而建议的答案给出正整数的有序序列. 通过替换而不是不替换来选择点的天真的修改确实允许完整的 non-negative 整数分布集，但它并没有使每个分布具有相同的可能性。考虑分配（4,3）：唯一的方法get [0, 0, 4] 就是滚动(0, 0)，但是你可以通过滚动(1, 3) 或(3, 1)).

得到[1, 2, 1]

Answer 2

如果没有大量的 (n x size x int32) space 开销，我找不到快速的 numpy 解决方案。所以这是@RichardYannow 的的 numba 实现。请记住从基准测试中排除第一个调用。

抽样算法是 np.rng.choice, size=1 samples should be "np.squeezed 的简化改编，以删除未使用的维度。

import numpy as np
import numba as nb # tested with numba 0.55.1

@nb.njit
def allocate(n, k, size):
    idx = np.zeros((size, k+1), np.int32)
    s = np.empty(n+k, np.uint8)
    for j in range(size):
        s.fill(0)
        for i in range(n+1, n + k):
            x = np.random.randint(i) + 1
            if s[x]: x = i
            s[x] = 1
            idx[j, i - n] = x

        idx[j, 1:k].sort()
        idx[j, k] = n+k
        for i in range(k):
            idx[j, i] = idx[j, i+1] - idx[j, i] - 1
    return idx[:, :k]

allocate(4, 3, size=5)

输出

array([[1, 1, 2],
       [0, 3, 1],
       [1, 1, 2],
       [1, 1, 2],
       [1, 2, 1]], dtype=int32)

检查所有可能分配的均匀分布

from collections import Counter
Counter(tuple(x) for x in allocate(4, 3, size=10000))

输出

Counter({(0, 0, 4): 685,
         (0, 1, 3): 680,
         (0, 2, 2): 646,
         (0, 3, 1): 666,
         (0, 4, 0): 662,
         (1, 0, 3): 660,
         (1, 1, 2): 670,
         (1, 2, 1): 661,
         (1, 3, 0): 647,
         (2, 0, 2): 650,
         (2, 1, 1): 699,
         (2, 2, 0): 682,
         (3, 0, 1): 667,
         (3, 1, 0): 669,
         (4, 0, 0): 656})

在 k 个容器中随机分配一个整数，并在分配中均匀分布

Allocate an integer randomly across k bins with uniform distribution across allocations

python

random