A more elegant way to uniformly sample from non-negative integers with a fixed sum

Question

Given a positive integer n, I would like to randomly output n non negative integers that sum to m (a positive integer).The output should be presented as a list of integers (not sorted).

My random sample should be uniformly sampled from all lists of n non negative integers that sum to m.

You can do this using a method inspired by stars and bars.

import random
def random_start(n, m):
    L = ([0]*m+[1]*(n-1))
    random.shuffle(L)
    return [len(item) for item in ''.join(map(str, L)).split("1")]

This works but fine but it is very ugly. How can I make this code more elegant?

Answer 1

Making your current implementation more readable with better naming. Your current code is reasonable, but the function and variable naming don't do much to assist with readability. Better names and a short comment or docstring can help quite a bit.

def random_segments(n_segments, total):
    '''
    Randomly divides TOTAL into segments of size 0 or larger.
    Returns a list of the segment lengths.
    '''
    SEG = '_'
    DIV = '|'
    segments = [SEG] * total
    dividers = [DIV] * (n_segments - 1)
    xs = segments + dividers
    random.shuffle(xs)
    return [len(seg) for seg in ''.join(xs).split(DIV)]

The primary drawbacks of your current implementation. (1) As noted in another review, your code does not scale well with respect to total: a huge value for that parameter causes your code to build a very large list of segments. (2) The logic moves from numbers to strings back to numbers: it would be more straightforward and natural to stay entirely in the realm of integers.

Converting your code's intuition into a better implementation: don't build any segments. To avoid those problems, just skip the segment-generation. Instead focus on computing the divider locations and the differences between adjacent dividers. At least to my eye, this version is both compact and intuitive. Of course, it won't win a speed race against numpy, so take that route if you need to handle very large n_segments.

from random import randint

def random_segments(n_segments, total):
    '''
    Randomly divides TOTAL into segments of size 0 or larger.
    Returns a list of the segment lengths.
    '''
    dividers = sorted(randint(0, total) for _ in range(n_segments - 1))
    z = zip([0] + dividers, dividers + [total])
    return [d2 - d1 for d1, d2 in z]

Answer 2

"Elegant" is an intangible, so let's specify:

• Can we make it faster? Yes (a lot!)
• Can we make it stay in the numeric domain instead of a weird venture into strings? Yes!
• Can we code-golf it down to less Python? Only a little, but no one should care. @lukstru demonstrates what this looks like, and whereas it's interesting academically, production code should not be code-golfed.

Aside from the fact that your approach invites string manipulation into an algorithm where it doesn't really fit, it introduces a runtime complexity that is polynomial in m, where a proper solution should be O(1) in m (i.e. not be slowed down whatsoever by high values of m).

You never replied about what typical range of values we would expect for n, so let's assume that it can get large. A vetorised Numpy solution could look like the following:

• Generate n uniformly-distributed floating-point values in [0, 1)
• Normalize by their sum
• Round and cast to int
• Calculate the rounding error as the difference between the actual and desired sums
• Compensate for the rounding error by incrementing or decrementing a fixed number of values by 1 as appropriate

Or, closer to your original strategy,

• Generate n - 1 randomised cuts
• Sort them
• Apply a numerical differential

Suggested

Containing comparative benchmarks.

import math
import random
from collections import defaultdict
from statistics import mean
from timeit import timeit

import numpy as np
from numpy.random import default_rng

rand = default_rng()


def op(n: int, m: int) -> list[int]:
    L = ([0]*m+[1]*(n-1))
    random.shuffle(L)
    return [len(item) for item in ''.join(map(str, L)).split("1")]


def fmc(n_segments, total):
    '''
    Randomly divides TOTAL into segments of size 0 or larger.
    Returns a list of the segment lengths.
    '''
    dividers = sorted(random.randint(0, total) for _ in range(n_segments - 1))
    z = zip([0] + dividers, dividers + [total])
    return [d2 - d1 for d1, d2 in z]


def lukstru_a(n: int, m: int) -> list[int]:
    cuts = sorted([math.floor(random.uniform(0, m)) for _ in range(n - 1)])
    values = []
    lastCut = 0
    for cut in cuts:
        values.append(cut - lastCut)
        lastCut = cut
    values.append(m - lastCut)
    return values


def lukstru_b(amount, totalSum) -> list[int]:
    cuts = [0] + sorted(random.choices(range(totalSum), k=amount - 1)) + [totalSum]
    return [cuts[index + 1] - cuts[index] for index in range(len(cuts) - 1)]


def numpy_sum(n: int, target_sum: int) -> np.ndarray:
    x = rand.random(size=n)
    x = (x * target_sum / x.sum()).round().astype(int)
    delta = x.sum() - target_sum
    if delta > 0:
        pool, = np.nonzero(x > 0)
        change = -1
    elif delta < 0:
        pool = n
        change = 1
    else:
        return x
    change_at = rand.choice(a=pool, size=abs(delta), replace=False)
    x[change_at] += change
    return x


def numpy_diff(n: int, m: int) -> np.ndarray:
    cuts = np.zeros(n + 1)
    cuts[1:-1] = rand.choice(a=m, size=n-1)
    cuts[-1] = m
    cuts.sort()
    return np.diff(cuts)


METHODS = (op, fmc, lukstru_a, lukstru_b, numpy_sum, numpy_diff)


def test() -> None:
    for method in METHODS:
        for n in (1, 5, 10):
            for target_sum in (10, 50, 100):
                result = method(n, target_sum)
                assert sum(result) == target_sum
                assert len(result) == n


def benchmark() -> None:
    times = defaultdict(list)
    for _ in range(10):
        for method in METHODS:
            def run():
                return method(10_000, 50_000)
            t = timeit(run, number=1)
            times[method.__name__].append(t)

    for method, method_times in times.items():
        print(f'{method}: {1e3 * mean(method_times):.1f} ms')


if __name__ == '__main__':
    test()
    benchmark()

Output

op: 58.6 ms
fmc: 19.6 ms
lukstru_a: 9.6 ms
lukstru_b: 7.3 ms
numpy_sum: 1.4 ms
numpy_diff: 1.3 ms

Uniformity testing

You ask:

It’s not clear to me however that your numpy code is sampling uniformly. Is there a simple proof?

For a test run of n=10, m=50 over 30,000 iterations we have

import numpy as np
from numpy.random import default_rng

rand = default_rng()


def rand_sum(n: int, target_sum: int) -> np.ndarray:
    x = rand.random(size=n)
    x = (x * target_sum / x.sum()).round().astype(int)
    delta = x.sum() - target_sum
    if delta > 0:
        pool, = np.nonzero(x > 0)
        change = -1
    elif delta < 0:
        pool = n
        change = 1
    else:
        return x
    change_at = rand.choice(a=pool, size=abs(delta), replace=False)
    x[change_at] += change
    return x


def test_uniform() -> None:
    total = np.zeros((10, 50), dtype=int)
    col_idx = np.arange(50)
    col_comparison = np.broadcast_to(col_idx, total.shape)

    for _ in range(30_000):
        x = rand_sum(n=10, target_sum=50)[:, np.newaxis]
        # total is 10 x 50.
        # Increment at a column equal to the value from x,
        # and at a row equal to the position in x.
        mask = col_comparison == x
        total[mask] += 1

    print(total[:, :20].T)


if __name__ == '__main__':
    test_uniform()

For each of the 10 output positions as columns, each row representing a value starting at 0, output frequencies are:

[[1354 1346 1404 1398 1380 1407 1385 1391 1395 1327]
 [2814 2776 2771 2808 2790 2761 2765 2768 2835 2806]
 [2959 3006 2944 2931 3004 3031 3046 2874 2899 2885]
 [2993 3052 2978 3088 3072 3033 3122 3048 3002 2997]
 [3208 3158 3111 3095 3161 3204 3165 3221 3149 3294]
 [3233 3255 3374 3299 3361 3380 3255 3380 3408 3448]
 [3502 3437 3483 3580 3499 3347 3517 3456 3539 3446]
 [3441 3510 3434 3392 3424 3434 3398 3463 3376 3494]
 [2931 2849 2903 2813 2788 2794 2825 2867 2907 2779]
 [1834 1865 1795 1844 1779 1920 1809 1802 1858 1858]
 [ 974  986  979  961  972  910  958  949  920  927]
 [ 432  440  474  429  438  410  464  435  385  422]
 [ 196  185  205  218  190  207  166  188  184  190]
 [  66   84   90   83   75   95   68  100   84   68]
 [  41   27   27   34   43   44   32   34   40   33]
 [  16   17   20   15   15   10   12   15   12   18]
 [   3    5    6    7    7    8    8    7    4    5]
 [   1    2    1    3    2    2    4    0    1    2]
 [   2    0    1    2    0    2    1    0    1    0]
 [   0    0    0    0    0    0    0    0    1    1]]

Across n, the distributions are equal, but across m the distribution is not uniform. However, it's not obvious that you actually care about this, since your original method exhibits much the same non-uniform distribution and is heavily biased toward small numbers; for n=3 m=50:

[[367 378 403]
 [400 365 381]
 [363 361 372]
 [353 350 356]
 [380 348 370]
 [349 342 362]
 [308 326 311]
 [333 357 354]
 [301 304 355]
 [317 315 324]
 [292 325 316]
 [302 317 276]
 [307 261 286]
 [298 290 283]
 [269 265 274]
 [254 283 317]
 [254 255 290]
 [242 259 263]
 [242 242 261]
 [245 233 244]
 [251 262 220]
 [204 215 199]
 [240 203 222]
 [216 216 203]
 [209 213 189]
 [189 186 197]
 [211 202 178]
 [183 184 193]
 [172 180 182]
 [187 178 167]
 [163 157 159]
 [163 168 161]
 [155 142 137]
 [134 151 124]
 [141 119 120]
 [113  92 114]
 [107 139 108]
 [118  96  97]
 [ 85 108  90]
 [ 81 101  84]
 [101  85  72]
 [ 89  80  71]
 [ 62  68  72]
 [ 57  62  52]
 [ 39  61  51]
 [ 48  46  46]
 [ 36  36  35]
 [ 30  16  26]
 [ 18  26  20]
 [ 15  21   9]]

Answer 3

This is longer, but more readable, and if you can reduce the loop with syntactic sugar we might even have an elegant solution

import random

def random_start(n, m):
    cuts = sorted([math.floor(random.uniform(0, m)) for _ in range(n - 1)])
    values = []
    lastCut = 0
    for cut in cuts:
        values.append(cut - lastCut)
        lastCut = cut
    values.append(m - lastCut)
    return values

EDIT: found a shorter version, it's less readable now though. Seems like you need some space to make the intent clear.

import random

def random_start2(amount, totalSum):
    cuts = [0] + sorted(random.choices(range(totalSum), k=amount - 1)) + [totalSum]
    return [cuts[index + 1] - cuts[index] for index in range(len(cuts) - 1)]