有效地选择具有给定概率的随机矩阵索引

时间:2016-09-23 22:32:08

标签: python numpy matrix

我有一系列概率,例如:

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

我想从该矩阵中选择一个元素(例如,选择一些索引(i,j)),并根据该矩阵对概率进行加权。这将使用的实际矩阵很大(高达1000x1000),所以我正在寻找一种有效的方法来做到这一点。这是我目前的解决方案:

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

这绝对没有效率。 (基本上,线性化矩阵,创建索引元组列表,然后相应地进行选择。)有没有更好的方法来进行此选择?

1 个答案:

答案 0 :(得分:3)

您不需要选择元组列表。只需使用arange(n)数组,然后按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))

结果:

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