Efficiently select random matrix indices with given probabilities
<p>I have a numpy array of probabilities, such as:</p>

<pre><code>[[0.1,  0,   0.3,],
  0.2,  0,   0.05],
  0,    0.15, 0.2 ]]
</code></pre>

<p>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:</p>

<pre><code>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
</code></pre>

<p>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?</p>

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

<pre><code>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))
</code></pre>

<p>The result:</p>

<pre><code>array([[0, 2],
       [2, 2],
       [2, 1],
       ..., 
       [1, 0],
       [0, 2],
       [0, 0]], dtype=int64)
</code></pre>

python numpy matrix