Блог

Рассмотрим следующий фрагмент кода:

 import numpy as np

mask = np.array([True, True, False, True, True, False])
val = np.array([9, 3])
arr = np.random.randint(1, 9, size = (5,len(mask)))
  

Как и ожидалось, мы получаем массив случайных целых чисел от 1 до 9 с 5 строками и 6 столбцами, как показано ниже. val Массив еще не использовался.

 [[2, 7, 6, 9, 7, 5],
 [7, 2, 9, 7, 8, 3],
 [9, 1, 3, 5, 7, 3],
 [5, 7, 4, 4, 5, 2],
 [7, 7, 9, 6, 9, 8]]
  

Теперь я представлю val = [9, 3] .
Где mask = True я хочу, чтобы элемент строки выбирался случайным образом от 1 до 9.
Где mask = False я хочу, чтобы элемент строки выбирался случайным образом от 1 до 3.

Как это можно сделать эффективно? Пример вывода показан ниже.

 [[2, 7, 2, 9, 7, 1],
 [7, 2, 1, 7, 8, 3],
 [9, 1, 3, 5, 7, 3],
 [5, 7, 1, 4, 5, 2],
 [7, 7, 2, 6, 9, 1]]
  

Одна из идей состоит в том, чтобы произвольно выбирать от 0 до 1, затем умножать на 9 или 3 в зависимости от mask и, наконец, добавлять 1 для перемещения выборки.

 rand = np.random.rand(5,len(mask))
is3 = (1-mask).astype(int)

# out is random from 0-8 or 0-2 depending on `is3`
out = (rand*val[is3]).astype(int)

# move out by `1`:
out = (out   1)
  

Вывод:

 array([[4, 9, 3, 6, 2, 1],
       [1, 8, 2, 7, 1, 3],
       [8, 2, 1, 2, 3, 2],
       [4, 3, 2, 2, 3, 2],
       [5, 8, 1, 5, 6, 1]])