Блог

Существует ли в Python функция, которая выполняет выборку из n-мерного массива numpy и возвращает индексы каждого розыгрыша. Если нет, то как можно было бы определить такую функцию?

Например.:

 >>> probabilities = np.array([[.1, .2, .1], [.05, .5, .05]])  
>>> print function(probabilities, draws = 10)
 ([1,1],[0,2],[1,1],[1,0],[0,1],[0,1],[1,1],[0,0],[1,1],[0,1])  
  

Я знаю, что эту проблему можно решить многими способами с помощью одномерных массивов. Однако я буду иметь дело с большими n-мерными массивами и не могу позволить себе изменять их форму только для выполнения одного рисования.

Вы можете использовать np.unravel_index :

 a = np.random.rand(3, 4, 5)
a /= a.sum()

def sample(a, n=1):
    a = np.asarray(a)
    choices = np.prod(a.shape)
    index = np.random.choice(choices, size=n, p=a.ravel())
    return np.unravel_index(index, dims=a.shape)

>>> sample(a, 4)
(array([2, 2, 0, 2]), array([0, 1, 3, 2]), array([2, 4, 2, 1]))
  

Это возвращает кортеж массивов, по одному на измерение a , длина каждого из которых равна количеству запрошенных выборок. Если вы предпочитаете массив формы (samples, dimensions) , измените оператор return на:

 return np.column_stack(np.unravel_index(index, dims=a.shape))
  

И теперь:

 >>> sample(a, 4)
array([[2, 0, 0],
       [2, 2, 4],
       [2, 0, 0],
       [1, 0, 4]])
  

Если ваш массив является непрерывным в памяти, вы можете изменить shape свой массив на месте:

 probabilities = np.array([[.1, .2, .1], [.05, .5, .05]]) 
nrow, ncol = probabilities.shape
idx = np.arange( nrow * ncol ) # create 1D index

probabilities.shape = ( 6, ) # this is OK because your array is contiguous in memory

samples = np.random.choice( idx, 10, p=probabilities ) # sample in 1D
rowIndex = samples / nrow # convert to 2D
colIndex = samples % ncol

array([2, 0, 1, 0, 2, 2, 2, 2, 2, 0])
array([1, 1, 2, 0, 1, 1, 1, 1, 1, 1])
  

Обратите внимание, что, поскольку ваш массив является непрерывным в памяти, reshape также возвращает представление:

 In [53]:

view = probabilities.reshape( 6, -1 )
view[ 0 ] = 9
probabilities[ 0, 0 ]
Out[53]:
9.0