2D версия numpy случайный выбор с взвешиванием

Question

Это относится к этому предыдущему сообщению: Numpy случайный выбор кортежей

У меня есть массив 2D numpy, и я хочу выбрать из него, используя массив вероятностей 2D. Единственный способ, которым я мог подумать, это сгладить, а затем использовать модуль и остаток для преобразования результата обратно в 2D-индекс

import numpy as np
# dummy data
x=np.arange(100).reshape(10,10)

# dummy probability array
p=np.zeros([10,10])
p[4:7,1:4]=1.0/9

xy=np.random.choice(x.flatten(),1,p=p.flatten())
index=[int(xy/10),(xy%10)[0]] # convert back to index
print(index)

, который дает

[5, 2]

но есть ли более чистый способ избежать сплющивания и по модулю? то есть я мог бы передать список кортежей координат как x, но как я могу тогда обработать веса?

Answer 1

Вы можете использовать numpy.random.randint для создания индекса, например:

# assumes p is a square array
ij = np.random.randint(p.shape[0], size=p.ndim) # size p.ndim = 2 generates 2 coords

# need to convert to tuple to index correctly
p[tuple(i for i in ij))]
>>> 0.0

Вы также можете индексировать несколько случайных значений одновременно:

ij = np.random.randint(p.shape[0], size=(p.ndim, 5)) # get 5 values
p[tuple(i for i in ij))]
>>> array([0.        , 0.        , 0.        , 0.11111111, 0.        ])

Answer 2

Я не думаю, что можно напрямую указать двумерный массив вероятностей. Так что бродить должно быть хорошо. Однако, чтобы получить соответствующие двумерные индексы из плоского индекса, вы можете использовать np.unravel_index

index= np.unravel_index(xy.item(), x.shape)
# (4, 2)

---

Для нескольких индексов вы можете просто сложить результат:

xy=np.random.choice(x.flatten(),3,p=p.flatten())

indices = np.unravel_index(xy, x.shape)
# (array([4, 4, 5], dtype=int64), array([1, 2, 3], dtype=int64))
np.c_[indices]
array([[4, 1],
       [4, 2],
       [5, 3]], dtype=int64)

, где np.c_ складывается вдоль правой оси и дает тот же результат, что и

np.column_stack(indices)