Эффективный способ выборки из индексов массива Numpy?

Я хотел бы попробовать из индексов массива 2D Numpy, учитывая, что каждый индекс взвешен по числу внутри этого массива. То, что я знаю, это с numpy.random.choice, однако это не возвращает индекс, а сам номер. Есть ли эффективный способ сделать это?

Вот мой код:

import numpy as np
A=np.arange(1,10).reshape(3,3)
A_flat=A.flatten()
d=np.random.choice(A_flat,size=10,p=A_flat/float(np.sum(A_flat)))
print d

Ответ 1

Вы можете сделать что-то вроде:

import numpy as np

def wc(weights):
    cs = np.cumsum(weights)
    idx = cs.searchsorted(np.random.random() * cs[-1], 'right')
    return np.unravel_index(idx, weights.shape)

Обратите внимание, что cumsum является самой медленной частью этого, поэтому, если вам нужно сделать это повторительно для одного и того же массива, я бы предложил вычислить cumsum раньше времени и повторно использовать его.

Ответ 2

Развернуть мой комментарий: Адаптация метода взвешенного выбора, представленного здесь fooobar.com/info/28990/...

def weighted_choice_indices(weights):
    cs = np.cumsum(weights.flatten())/np.sum(weights)
    idx = np.sum(cs < np.random.rand())
    return np.unravel_index(idx, weights.shape)