Numpy случайный выбор с ограниченным временем повторения для каждого элемента

Question

Дано N элементов 0, ..., N-1 и массива repeats размером N, которые определяют количество каждого элемента. Интересно, какой самый быстрый способ случайным образом выбрать size элементов? Результатом является массив размером N, указывающий выбранное количество каждого элемента. sum(output) == size.

Например,

def choice(repeats, size): ...

choice([100, 200, 300, 400, 500], 5)
# a random output: array([1, 0, 2, 1, 1])

См. Ответ для текущего решения.

Интересно, есть ли способ избежать выделения option массив из sum(repeats) элементов?

По сравнению с неправильной реализацией, но без выделения, разница в скорости составляет около 20 раз при большом количестве элементов:

def choice_wrong(repeats, size):
    chosen = np.random.choice(len(repeats), size, p=repeats/np.sum(repeats))
    return np.bincount(chosen, minlength=len(repeats))

%timeit choice([10000, 20000, 30000, 40000, 50000], 5000)
# 2.3 ms ± 7.47 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit choice_wrong([10000, 20000, 30000, 40000, 50000], 5000)
# 129 µs ± 822 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Answer 1

Текущее решение, которое у меня есть:

def choice(repeats, size):
    assert size <= np.sum(repeats)
    options = np.concatenate([np.repeat(i, r) for i, r in enumerate(repeats)])
    chosen = np.random.choice(options, size=size, replace=False)
    return np.bincount(chosen, minlength=len(repeats))