NumPy 中 random.choice 的多维数组答案

【问题标题】：Multidimensional array for random.choice in NumPyNumPy 中 random.choice 的多维数组
【发布时间】：2017-06-18 08:55:56
【问题描述】：

我有一张桌子，我需要使用 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) 中的元素具有相同的概率，因此我将看到来自 arr0 等的任何元素的概率为 0.5。

有什么优雅的方法吗？

【问题讨论】：

这看起来像 3 1D 数组而不是 3D 数组。你能确认一下吗？
是的，感谢您的更正。基本上，我最终想要的是 >>>np.random.choice([t, pos, n], 5, p=C1)

标签： python arrays numpy multidimensional-array

【解决方案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)

【讨论】：

【解决方案2】：

这是我带来的。它采用概率向量或权重按列组织的矩阵。权重将归一化为总和为 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])

【讨论】：