Numpy: случайным образом выбрать 1 выборку из каждого из N дискретных распределений вероятности

Question

У меня есть двумерный массив Numpy P, где каждый P[ii,:] - это дискретное распределение вероятностей по 0,1, ..., M-1. Я хотел бы попробовать каждый из P[ii,:] один раз. Я пытался использовать функцию numpy.random.choice:

for ii in range(N):
    samples[ii] = np.random.choice(np.arange(M), p=P[ii,:])

Как я могу сделать то же самое без цикла for?

Спасибо.

Answer 1

Я думаю, что вы, вероятно, хотите np.apply_along_axis

import numpy as np

P = np.array([[0,1/2,1/2],[1,0,0],[1/3,1/3,1/3]])

def rand_choice(c):
  return np.random.choice(c.size, p=c)

samples = np.apply_along_axis(rand_choice, 1, P)