Создать массив numpy на основе нескольких условий для двух массивов numpy

Question

Учитывая два массива как таковых

a = np.array([0,1,0,1,0,0])
b = np.array([0,0,1,1,1,0])

Я хочу создать массив numpy на основе условий как

output = np.zeros(len(arr1))

for i in range(0, len(arr1)):
    if a[i] == 0 and b[i] == 0:
        output[i] = 0
    elif a[i] == 1 and b[i] == 0:
        output[i] = 1
    elif a[i] == 0 and b[i] == 1:
        output[i] = 2
    elif a[i] == 1 and b[i] == 1:
        output[i] = 3


expected_output = np.array([0,1,2,3])

Какой самый быстрый способ создания такого массива?

Answer 1

Остальные привели примеры, как это сделать на чистом python. Если вы хотите сделать это с массивами с 100.000 элементов, вы должны использовать numpy:

In [1]: import numpy as np
In [2]: vector1 = np.array([0,1,0,1,0,1])
In [3]: vector2 = np.array([0,0,1,1,2,2])

Делать поэлементное добавление теперь так же тривиально, как и

In [4]: sum_vector = vector1 + vector2 * 2
In [5]: print(sum_vector) # python3.x kaugh...
[0, 1, 2, 3, 4, 5]

, как в Matlab; -)

U также может видеть: Поэлементное добавление 2 списков?

Answer 2

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

Это:

for i in range(0, len(arr1)):
    if a[i] == 0 and b[i] == 0:
        output[i] = 0
    elif a[i] == 1 and b[i] == 0:
        output[i] = 1
    elif a[i] == 0 and b[i] == 1:
        output[i] = 2
    elif a[i] == 1 and b[i] == 1:
        output[i] = 3

Эквивалентно:

output[a == 0 & b == 0] = 0
output[a == 1 & b == 0] = 1
output[a == 0 & b == 1] = 2
output[a == 1 & b == 1] = 3

или:

a0 = a == 0
b0 = b == 0
a1 = a == 1
b1 = b == 1
output[a0 & b0] = 0
output[a1 & b0] = 1
output[a0 & b1] = 2
output[a1 & b1] = 3

Answer 3

Как указано выше, sugested, это самый быстрый \ простой способ:

import numpy as np

a = np.array([0,1,0,1,0,0])
b = np.array([0,0,1,1,1,0])

a + 2*b