Нахождение случайного вхождения в массиве NumPy, когда более одного выбора

Question

Я написал этот код для достижения цели возврата случайного значения из списка элементов, соответствующих условию предиката:

N=<int>
sampl = np.random.randint(low=0, high=N+1, size=(10,))
xs = np.where(sampl == 1)
ys = np.array([tuple(x) for x in xs], dtype=int)[0]
x = np.random.choice(ys)

Пример: если я запускаю код с N=2 и ищу только 1 с в массиве:

    sampl = np.random.randint(low=0, high=N+1, size=(10,))

--> sampl = [2 1 0 0 0 1 0 0 2 1]

    xs = np.where(sampl == 1)

--> [2 1 0 0 0 1 0 0 2 1]  # Positions 1, 5, 9 are of interest. 
       ^       ^       ^ 

    ys = np.array([tuple(x) for x in xs], dtype=int)[0]

--> ys = [1 5 9] # Put them in an array. 

    x = np.random.choice(ys)

--> x = 9 # Pick a random one and return it

Это работает, но это не лаконично, и я столкнулся с несколькими проблемами, пытаясь сделать его более элегантным.

• numpy.where() возвращает кортеж, когда не проходит ничего, кроме условия. Я попытался передать x=sampl, но среда выполнения жалуется, говоря, что функция не принимает параметры (это происходит, когда я проверяю код).
• Опять же, создание массива из кортежа заставляет меня возвращать первый элемент. Это подвержено ошибкам при тестировании на крайние случаи (например, предикат не обнаружил значения).

У вас есть предложения по улучшению этого кода? Я хочу придерживаться numpy / pandas, так как массивы станут очень большими.

Answer 1

Вероятно, самый изящный способ, который я мог придумать, - это произвольно перемешать ваш массив, а затем выполнить первое вхождение. Это должно быть довольно кратким.

Так что-то вроде:

np.random.shuffle(sampl)
x = np.ravel(np.where(sampl==1))[0]

или, как вы предложили, без перетасовки, это выглядело бы как

x = np.random.choice(np.ravel(np.where(sampl==1)))

Если подумать, я думаю, что метод choice будет бесконечно быстрее, чем тасование.

Следующая проблема - крайние случаи. Как справиться с этим, зависит от того, какое поведение по умолчанию вы ожидаете. Если вы ожидаете, что в большинстве случаев условие повлечет за собой хотя бы одно попадание, вам следует обработать случай, когда нет попадания с исключением:

try: 
   x = np.random.choice(np.ravel(np.where(sampl==1)))
except: 
   # TODO
   pass

Я бы настоятельно рекомендовал сделать это, если только вы редко найдете хит. Но не верь мне на слово ... время сам.

Другим вариантом было бы поставить условие, которое явно проверяет, что

np.size( np.where(sampl==1) ) > 0

прежде чем продолжить. Однако я бы предположил, что этот подход медленнее, чем try...except.