Generate random array of integers with known total

Here is an exact (every legal sum has same probability) solution. It uses an enumeration of all legal sums, not in the sense that we go through each and every sum, but rather given a number n we can directly compute the n-th sum in the enumeration. As we also know the total number of legal sums, we can simply draw uniform integers and transform them:

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

Demo:

# 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])

Brief explanation:

We use a simple recurrence to compute the total number of capped partitions. We split on the first bin, i.e we fix the number in the first bin and by recurrence retrieve the number of ways to fill the remaining bins. Then we simply sum over the different first bin options. We use a cache decorator to keep the recursion under control. This decorator remembers all parameter combinations we've already done, so when we arrive at the same parameters via different recursion paths we don't have to do the same computation again.

The enumeration works similarly. Assume lexicographic order. How to find the n-th partition? Again, split on the first bin. Because we know for each value the first bin can take the total number of ways the remaining bins can be populated we can form the accumulated sum and then see where n fits in. If n lies between the i-th and i+1st partial sum the the first index is i+low, we subtract the i-th sum from n and start over on the remaining bins.

Here is a quick & dirty version (not easily vectorizable and with no guarantees of any particular distribution over possible outputs):

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))

It chooses random integers one at a time.

At each point it knows how much of the total needs to be (randomly) distributed between how many remaining numbers (num). It uses this knowledge to adjust the low and high bounds to remain within the constraints throughout.

This method biases numbers that are chosen later towards having larger values. The shuffle removes that bias (but, again, does not provide any guarantees about the overall distribution over possible outputs).

I'd be delighted to hear if there's a fundamentally better way to do this.

Modified your code to support array,

>>> 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)

Output:

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

you can do it like below using multinomial where np.random.multinomial(s, [1/5]*5) is similar to drawing an unbiased dice with 5 sides s times and it returns number of times each side appeared on each draw. so sum of the counts each appeared must be equal to number of total draws

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)

The following differs from my other answer in that it guarantees uniform distribution over all possible outputs, but is probably only practical for small inputs (where the space of possible outputs is fairly small):

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))

This materialises all possible outputs into a list before picking one. In practice, one would probably use reservoir sampling (of size 1).

If you need to generate multiple random sets with the same parameters, reservoir sampling with replacement might be the answer: https://epubs.siam.org/doi/pdf/10.1137/1.9781611972740.53

Here is a maybe dirtier solution than others. But it will give uniform distribution as long as random.sample gives it.

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]

It starts with a list of ones, keeps adding +1 until it reaches the goal, also keeps track of which value reached 5, so no longer adds there. Pretty basic and slow solution. (Still O(n))

Edit: I have written it for a single value, you would have to apply it in a loop for multiple results.