Create numpy array based on multiple conditions on two numpy arrays

Question

Given two arrays as such

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

I want to create a numpy array based on conditions as

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])

what is the fastest way of creating such array?

edited. I hope it's clearer now. I have tried using np.where, but it doesn't seem to support different outputs from different conditions. — Hackore, Jan 22 '20 at 00:47
Err... Isn't this just counting upwards? This isn't very clear still. — Mateen Ulhaq, Jan 22 '20 at 00:54
Looking at the older examples you gave, I'm curious, what is this for? — AMC, Jan 22 '20 at 03:08

score 1 · Accepted Answer · answered Jan 22 '20 at 01:02

The others gave examples how to do this in pure python. If you want to do this with arrays with 100.000 elements, you should use 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])

Doing the element-wise addition is now as trivial as

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

just like in Matlab ;-)

U may see also : Element-wise addition of 2 lists?

score 0 · Answer 2 · answered Jan 22 '20 at 01:16

0

As answer above sugested, this is the fast\easiest way:

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

answered Jan 22 '20 at 01:16

Oleg Kmechak

165
1
14

score 0 · Answer 3 · answered Jan 22 '20 at 01:59

In the most generalized sense, you can use boolean arrays to assign values to indices matching whatever criteria that you want.

This:

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

Is equivalent to:

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

or:

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

Create numpy array based on multiple conditions on two numpy arrays

3 Answers3