Самый быстрый способ поиска значений в массиве, которые также находятся в другом массиве

Question

У меня есть два массива: arrayA (3D), который мне нужно проанализировать, и arrayB, в которой хранятся значения интереса (2D). Я хочу вернуть третий массив (те же размеры, что и у исходного массива A), который имеет 1 с, где значение является одним из интересующих значений (в массиве B), или 0, если нет.

Это код, который я Я использую:

arrayC = np.zeros((arrayA.shape[0], arrayA.shape[1], arrayA.shape[2]))
for k in range(arrayA.shape[2]):
    for i in range(arrayA.shape[0]):
        for j in range(arrayA.shape[1]):
            if arrayA[i][j][k] in arrayB[k]:
                arrayC[i][j][k] = 1 

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

(я тоже пробовал с np.where, но не могу заставить его работать с условием диапазона)

Answer 1

Вы можете получить ожидаемый результат в векторизации с помощью np.isin и np.where:

arrayC = np.where([np.isin(arrayA[i], arrayB[i])
                   for i in range(len(arrayB))], 1, 0)

Answer 2

Вы можете использовать numpy .isin :

c = np.zeros_like(a)
for k in range(a.shape[2]):
    c[:, :, k] = np.isin(a[:, :, k], b[k])

В грубом тесте на моей машине для a с формой (1000, 1000, 10) и b с формой (1000, 500) это занимает 1,55 секунды, а ваша версия - 50,3 секунды.