how to get the index of numpy.random.choice? - python

Question

Is it possible to modify the numpy.random.choice function in order to make it return the index of the chosen element? Basically, I want to create a list and select elements ran

Answer 1

Maybe late but it worth to mention this solution because I think the simplest way to do so is:

a = [1, 4, 1, 3, 3, 2, 1, 4]
n = len(a)
idx = np.random.choice(list(range(n)), p=np.ones(n)/n)

It means you are choosing from the indices uniformly. In a more general case, you can do a weighted sampling (and return the index) in this way:

probs = [.3, .4, .2, 0, .1]
n = len(a)
idx = np.random.choice(list(range(n)), p=probs)

If you try to do so for so many times (e.g. 1e5), the histogram of the chosen indices would be like [0.30126 0.39817 0.19986 0. 0.10071] in this case which is correct.

Anyway, you should choose from the indices and use the values (if you need) as their probabilities.