Given N elements 0, ..., N-1, and an array repeats of size N which specify the quantity of each element. I wonder what's the fastest way to randomly choose size elements? The output is an array of size N indicating the chosen quantities of each element. sum(output) == size.
For example,
def choice(repeats, size): ...
choice([100, 200, 300, 400, 500], 5)
# a random output: array([1, 0, 2, 1, 1])
See the answer for a current solution.
I wonder if there is anyway to avoid the allocation of the option array of sum(repeats) elements?
Comparing to an incorrect implementation but without the allocation, there are around 20x speed difference when there are many elements:
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)
Aucun commentaire:
Enregistrer un commentaire