I have a numpy array P with size (N,M), where each row is a discrete probability distribution over 0, ..., M -1 (i.e. each row sums to one and all elements are non-negative). I also have an array x with length K containing elements in 0, ..., N-1. For each x[k], my goal is to generate one sample according to distribution P[x[k], :] and store them in array y with length K.
I have written a code for this task with a loop
y = np.zeros(K)
for i, k in enumerate(K):
y[i] = np.random.choice(np.arange(M), p = P[x[k],:])
Is it possible to compute y without using a loop?