Search code examples
pythonif-statementlist-comprehensionenumeration

Python Enumeration Using List Comprehension and If Statement


Is there a way to enumerate in list comprehension where the enumeration only increments when the if statement is True?

For example, the code below enumerates whether the if statement is True or False and I would like to have a continuous index.

my_list = [0,1,2] [i for i, w in enumerate(my_list) if w != 1]

I also tried the following, but you can't use pass inside a list comprehension like below.

[i if w != 1 else pass for i, w in enumerate(my_list)]

Solution

  • There is. One simple way is to use itertools.count:

    >>> from itertools import count
    >>> c = count(0)
    >>> my_list = [0,1,2]
    >>> [next(c) for w in my_list if w != 1]
    [0, 1]
    

    HOWEVER: This involves state-change inside a list-comprehension, which is just... not good. Note, if you run your comprehension again, this will have unfortunate effects:

    >>> [next(c) for w in my_list if w != 1]
    [2, 3]
    >>> [next(c) for w in my_list if w != 1]
    [4, 5]
    

    Instead, I would do something like:

    >>> range(sum(w != 1 for w in my_list))
    range(0, 2)