Search code examples
python-3.xlistlist-comprehension

Getting indices using conditional list comprehension


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


Solution

  • 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