Efficiently select random matrix indices with given probabilities

Question

I have a numpy array of probabilities, such as:

[[0.1,  0,   0.3,],
  0.2,  0,   0.05],
  0,    0.15, 0.2 ]]

I want to select an element (e.g., select some indices (i,j)) from this matrix, with probability weighted according to this matrix. The actual matrices this will be working with are large (up to 1000x1000), so I'm looking for an efficient way to do this. This is my current solution:

def weighted_mat_choice(prob_mat):
    """
    Randomly select indices of the matrix according to the probabilities in prob_mat
    :param prob_mat: Normalized probabilities to select each element
    :return: indices (i, j) selected
    """
    inds_mat = [[(i, j) for j in xrange(prob_mat.shape[1])] for i in xrange(prob_mat.shape[0])]
    inds_list = [item for sublist in inds_mat for item in sublist]
    inds_of_inds = xrange(len(inds_list))
    prob_list = prob_mat.flatten()
    pick_ind_of_ind = np.random.choice(inds_of_inds, p=prob_list)
    pick_ind = inds_list[pick_ind_of_ind]
    return pick_ind

which is definitely not efficient. (Basically, linearizing the matrix, creating a list of index tuples, and then picking accordingly.) Is there a better way to do this selection?

Is your matrix of probabilities sparse? That is, are most of the values 0? — Warren Weckesser
No. Actually, no elements of the matrix in my case are ever 0. — Julia Ebert

HYRY HYRY · Accepted Answer · 2016-09-23T22:38:50

You don't need a list of tuple to choice. Just use a arange(n) array, and convert it back to two dimension by unravel_index().

import numpy as np

p = np.array(
[[0.1,  0,   0.3,],
 [0.2,  0,   0.05],
 [0,    0.15, 0.2]]
)

p_flat = p.ravel()
ind = np.arange(len(p_flat))

res = np.column_stack(
    np.unravel_index(
        np.random.choice(ind, p=p_flat, size=10000),
        p.shape))

The result:

array([[0, 2],
       [2, 2],
       [2, 1],
       ..., 
       [1, 0],
       [0, 2],
       [0, 0]], dtype=int64)

Efficiently select random matrix indices with given probabilities

1 Answers