Генерация случайного массива целых чисел с известным итогом

Question

Рассмотрим следующий игрушечный массив с целыми числами в диапазоне от 5 - 25:

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

Как создать массивы 5 случайных целых чисел в диапазоне от 1 - 5,сумма которых равна каждому числу в массиве a? Решение должно генерировать равномерное распределение по всем возможным выходам.

До сих пор мне удалось создать простую функцию, которая будет выдавать 5 случайные целые числа:

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. Мы используем декоратор кеша, чтобы контролировать рекурсию. Этот декоратор запоминает все комбинации параметров, которые мы уже сделали, поэтому, когда мы приходим к одним и тем же параметрам по разным путям рекурсии, нам не нужно снова делать одни и те же вычисления.

Перечисление работает аналогично. Предположим, лексикографический порядок. Как найти n-й раздел? Опять же, разделить на первую корзину. Поскольку мы знаем, что для каждого значения первый бин может принимать общее количество способов заполнения оставшихся бинов, мы можем сформировать накопленную сумму и затем посмотреть, куда вписывается n. Если n лежит между i-й и i + 1-й частичной суммойЕсли первый индекс равен i + low, мы вычитаем i-ую сумму из n и начинаем заново с оставшихся бинов.

Answer 2

Вот, пожалуй, более грязное решение, чем другие. Но это даст равномерное распределение, пока 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 3

Следующее отличается от моего другого ответа тем, что оно гарантирует равномерное распределение по всем возможным выходам, но, вероятно, практично только для небольших входов (где пространство возможных выходов довольно мало):

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

Это материализует все возможные результаты в список, прежде чем выбрать один. На практике, вероятно, можно использовать отбор проб из резервуара (размера 1).

Если вам нужно создать несколько случайных наборов с одинаковыми параметрами, отбор проб из резервуара с заменой может быть ответом: https://epubs.siam.org/doi/pdf/10.1137/1.9781611972740.53

Answer 4

вы можете сделать это, как показано ниже, используя многочлен, где 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 5

Изменен ваш код для поддержки массива,

>>> 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 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 устраняет это смещение (но, опять же, не дает никаких гарантий относительно общего распределения по возможным выходам).

Я был бы рад услышать, если есть принципиально лучший способ сделать это.