Python: использовать массив в качестве условия для другого массива

Question

Можно ли просто использовать массив вроде:

import numpy as np
a = np.array([1,76,120,158,199]

как массив условий для другого массива?

У меня есть, например, массив вроде:

b=np.array([0,1,2,3,4,5,76,77,78,79,120,121,122,158,159,199,200,201])

Все числа массива b, которые также находятся в массиве a, должны получить число 1, например:

возврат:

c=np.array([0,1,2,3,4,5,1,77,78,79,1,121,122,1,159,1,200,201])

Так есть быстрая версия для этого. Потому что фактический массив - это огромный трехмерный массив с множеством чисел, который, конечно, можно сплющить.

Answer 1

Использование numpy.in1d и логическое индексирование:

b[np.in1d(b, a)] = 1

print(b)

array([  0,   1,   2,   3,   4,   5,   1,  77,  78,  79,   1, 121, 122,
         1, 159,   1, 200, 201])