Answers already given vectorize by over generating and throwing up some outputs, it seems wrong.
Moreover, I will generalize to any number of dices.
First, you need to be able to get a condlist: it is a list of length the number of dices, with each i-th element being a boolean array containing True where the i-th dice should be used:
dices_idxs = np.array([0, 1, 2])
dices_sequence = np.array([0, 1, 2, 2, 1, 1, 0])
condlist = np.equal(*np.broadcast_arrays(dices_sequence[None, :], dices_idxs[:, None]))
print(condlist)
# [[ True False False False False False True]
# [False True False False True True False]
# [False False True True False False False]]
Second, you can generalize the answer given by @Ahmed AEK using np.select:
def dice_simulator_select(dices_sequence, dices_weights):
faces = np.arange(1, 7)
num_dices = len(dices_weights)
dices_idxs = np.arange(num_dices)
num_throws = len(dices_sequence)
condlist = list(
np.equal(*np.broadcast_arrays(dices_sequence[None, :], dices_idxs[:, None]))
)
choicelist = [
RNG.choice(faces, size=num_throws, p=dices_weights[dice_idx])
for dice_idx in range(num_dices)
]
return np.select(condlist, choicelist)
But it has the issue stated first as it over-generates then discards some generated values, which can be problematic considering randomness.
A more correct way is to use np.piecewise:
def dice_simulator_piecewise(dices_sequence, dices_weights):
faces = np.arange(1, 7)
num_dices = len(dices_weights)
dices_idxs = np.arange(num_dices)
num_dices = len(dices_weights)
condlist = list(
np.equal(*np.broadcast_arrays(dices_sequence[None, :], dices_idxs[:, None]))
)
# note size=len(x) ensure no more sample than needed are generated
funclist = [
lambda x: RNG.choice(faces, size=len(x), p=dices_weights[int(x[0])])
] * num_dices
return np.piecewise(dices_sequence, condlist, funclist)
You can use the functions as follows, and see that the correct function using np.piecewise is even faster (20% faster in below case):
RNG = np.random.default_rng()
dices_weights = [
None, # uniform
[1 / 12, 1 / 12, 1 / 12, 1 / 4, 1 / 4, 1 / 4],
None,
[1 / 4, 1 / 4, 1 / 4, 1 / 12, 1 / 12, 1 / 12],
None,
[1 / 12, 1 / 12, 1 / 12, 1 / 4, 1 / 4, 1 / 4],
]
num_dices = len(dices_weights)
num_throws = 1_000
dices_sequence = RNG.choice(np.arange(num_dices), size=num_throws)
%timeit dice_simulator_select(dices_sequence, dices_weights)
%timeit dice_simulator_piecewise(dices_sequence, dices_weights)
# 311 µs ± 5.94 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
# 240 µs ± 10.3 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)