I have a non-negative normalized vector p. I would like to sample an index from the index set of the vector. The probability getting sample k is p[k]. Using np.random.choice function, I can write the following code.

p = [0.2, 0.3, 0.1, 0.3, 0.1]
indices = np.arange(len(p))
k = np.random.choice(indices, p=p)

My question is, how can I generalize this code for multi-dimensional arrays? For example, given three dimensional non-negative normalized IxJxK tensor p = np.random.rand(I,J,K) how can I sample the index (i,j,k) with the probability p[i,j,k]?

I have a non-negative normalized vector p. I would like to sample an index from the index set of the vector. The probability getting sample k is p[k]. Using np.random.choice function, I can write the following code.

p = [0.2, 0.3, 0.1, 0.3, 0.1]
indices = np.arange(len(p))
k = np.random.choice(indices, p=p)

My question is, how can I generalize this code for multi-dimensional arrays? For example, given three dimensional non-negative normalized IxJxK tensor p = np.random.rand(I,J,K) how can I sample the index (i,j,k) with the probability p[i,j,k]?

• python
• probability
• sampling

1 Answer 1

Suppose that x is a probability matrix:

x = np.random.random((3, 4, 5))
x /= np.sum(x)

You can use flatten(x) and np.random.choice() to get values in the range 0..3x4x5 with the associated probabilities:

flat = x.flatten()
temp = np.random.choice(np.arange(len(flat)), 10, p=flat)
print(temp)

You can now convert the indices in temp into indices in the original array:

np.unravel_index(temp, x.shape))

Note that this will return a tuple of three numpy arrays, with the first array being the first index, the second element being the second index, and so on.