Getting indices using conditional list comprehension

Question

I have the following np.array:

my_array=np.array([False, False, False, True, True, True, False, True, False, False, False, False, True])

How can I make a list using list comprehension of the indices corresponding to the True elements. In this case the output I'm looking for would be [3,4,5,7,12]

I've tried the following:

cols = [index if feature_condition==True for index, feature_condition in enumerate(my_array)]

But is not working

KArrow'sBest · Accepted Answer · 2022-05-12T13:07:33.743

1

why specifically a list comprehension?

>>> np.where(my_array==True)
(array([ 3,  4,  5,  7, 12]),)

this does the job and is quicker. The list solution would be:

>>> [index for index, feature_condition in enumerate(my_array) if feature_condition == True]
[3, 4, 5, 7, 12]

Accepted answer of this explains the confusion of the ordering: if/else in a list comprehension

I was curious of the differences:

def np_time(array):
     np.where(my_array==True)
def list_time(array):
    [index for index, feature_condition in enumerate(my_array) if feature_condition == True]

timeit.timeit(lambda: list_time(my_array),number = 1000)
0.007574789000500459
timeit.timeit(lambda: np_time(my_array),number = 1000)
0.0010812399996211752

edited May 12 '22 at 13:07

answered May 12 '22 at 12:35

KArrow'sBest

150
9

Thanks! This is perfect :) Do you know what is wrong with my code though? – Caterina May 12 '22 at 12:38
@Caterina Your solution was on the right track - you just can't have the `if feature_condition` in the middle of the list comprehension. – ssp May 12 '22 at 13:18
@ssp I am a kind of newbie in py and to me it was very confusing the difference in the if/else case and the normal if case, always getting them swapped – KArrow'sBest May 12 '22 at 13:20

score 1 · Answer 2 · answered May 12 '22 at 12:39

The order of if is not correct, should be in last -

$more numpy_1.py
import numpy as np
my_array=np.array([False, False, False, True, True, True, False, True, False, False, False, False, True])
print (my_array)

cols = [index for index, feature_condition in enumerate(my_array) if feature_condition]

print (cols)

$python numpy_1.py
[False False False  True  True  True False  True False False False False
  True]
[3, 4, 5, 7, 12]

Getting indices using conditional list comprehension

2 Answers2