Python array simple issue

Question

I've the following two arrays:

array1 = [0, 1, 1, 0]
array2 = ['foo', 'bar', 'hello', 'bye']

I want to save into an array the values of array2 that has the index 1in array1.

On the example above, the desired result should be result_array = ['bar', 'hello'].

I've tried something like this, but it's not working.

for i in array1:
  if i = 1:
     result_array.append(array2[i])

Thanks in advance

Answer 1

The problem with your code is that you're using = in if condition, replace it with ==. And secondly to get the index as well as item you need to use enumerate, currently you're appending array[i], so your code will end up appending 'bar' two times.

>>> result_array = []
>>> for i, x in enumerate(array1):
        if x == 1:
            result_array.append(array2[i])
...         
>>> result_array
['bar', 'hello']

Another better way to do this is to use zip and a list comprehension:

>>> [b for a, b in zip(array1, array2) if a==1]
['bar', 'hello']

And fastest way to do this is to use itertools.compress:

>>> from itertools import compress
>>> list(compress(array2, array1))
['bar', 'hello']