NumPy中random.choice的多维数组

时间:2017-06-18 08:55:56

标签: python arrays numpy multidimensional-array

我有一张桌子,我需要使用random.choice进行概率计算, 例如(取自文档):

>>> aa_milne_arr = ['pooh', 'rabbit', 'piglet', 'Christopher']
>>> np.random.choice(aa_milne_arr, 5, p=[0.5, 0.1, 0.1, 0.3])
array(['pooh', 'pooh', 'pooh', 'Christopher', 'piglet'],
      dtype='|S11')

如果我有3D数组而不是aa_milne_arr,它不会让我继续。我需要为3个数组生成具有不同概率的随机事物,但对于它们内部的元素则相同。例如,

>>> arr0 = ['red', 'green', 'blue']
>>> arr1 = ['light', 'wind', 'sky']
>>> arr3 = ['chicken', 'wolf', 'dog']
>>> p = [0.5, 0.1, 0.4]

我想要arr0(0.5),arr1(0.1)和arr3(0.4)中元素的相同prob,因此我将看到0.5的概率,来自arr0等的任何元素。

这是一种优雅的方式吗?

2 个答案:

答案 0 :(得分:1)

p的值除以数组的长度,然后重复相同的长度。

然后从具有新概率的连续数组中选择

arr = [arr0, arr1, arr3]
lens = [len(a) for a in arr]
p = [.5, .1, .4]
new_arr = np.concatenate(arr)
new_p = np.repeat(np.divide(p, lens), lens)

np.random.choice(new_arr, p=new_p)

答案 1 :(得分:0)

这是我的意思。 它采用概率向量,或者权重按列组织的矩阵。权重将归一化为1。

import numpy as np

def choice_vect(source,weights):
    # Draw N choices, each picked among K options
    # source: K x N ndarray
    # weights: K x N ndarray or K vector of probabilities
    weights = np.atleast_2d(weights)
    source = np.atleast_2d(source)
    N = source.shape[1]
    if weights.shape[0] == 1:
        weights = np.tile(weights.transpose(),(1,N))
    cum_weights = weights.cumsum(axis=0) / np.sum(weights,axis=0)
    unif_draws = np.random.rand(1,N)
    choices = (unif_draws < cum_weights)
    bool_indices = choices > np.vstack( (np.zeros((1,N),dtype='bool'),choices ))[0:-1,:]
    return source[bool_indices]

它避免使用循环,就像random.choice的矢量化版本。

然后你可以这样使用它:

source = [[1,2],[3,4],[5,6]]
weights = [0.5, 0.4, 0.1]
choice_vect(source,weights)
>> array([3, 2])

weights = [[0.5,0.1],[0.4,0.4],[0.1,0.5]]
choice_vect(source,weights)
>> array([1, 4])