find first element and index matching condition in list

Viewed 2614

Consider this simple example

mylist = [-1,-2,3,4,5,6]

for idx, el in enumerate(mylist):
    if el > 0:
        myidx, myel = idx, el
        break

myidx, myel
Out[20]: (2, 3)

I am interested in finding the first index and the corresponding first element in a python list that matches a specific condition (here, this is simply > 0).

In the code above, I loop over the elements using enumerate and then use the if clause to find the correct elements. This looks very cumbersome to me. Is there a better way to do this? Using a native python function for instance?

Thanks!

3 Answers

Something like this should work:

l = [-1,-2,3,4,5,6]
list(x > 0 for x in l).index(True)
# Output: 2

To find all patters, we can use python built in functions using

from itertools import filterfalse
f = filterfalse(lambda x: x[1] <= 0, enumerate(l))
print(list(f))
# [(2, 1), (3, 2), (4, 3)]

You could do it in a list comprehension. This is basically the same as your code but condensed into one line, and it builds a list of results that match the criteria.

The first way gets all the matches

mylist = [-1,-2,3,4,5,6]

results = [(i, el) for i, el in enumerate(mylist) if el > 0]

Another way would be to use a generator expression which is probably faster, and just unpack it. This gets the first one.

*next((i, el) for i, el in enumerate(mylist) if el > 0))

This loops the list and checks the condition, then puts the index and element into a tuple. Doing this inside parentheses turns it into a generator, which is much faster because it hasn't actually got to hold everything in memory, it just generates the responses as you need them. Using next() you can iterate through them. As we only use next() once here it just generates the first match. Then we unpack it with *

As there are two other valid answers here I decided to use timeit module to time each of them and post the results. For clarity I also timed the OP's method. Here is what I found:

import timeit
# Method 1 Generator Expression
print(timeit.timeit('next((i, el) for i, el in enumerate([-1,-2,3,4,5,6]) if el > 0)', number=100000))
0.007089499999999999

# Method 2 Getting index of True
print(timeit.timeit('list(x > 0 for x in [-1,-2,3,4,5,6]).index(True)', number=100000))
0.008104599999999997

# Method 3 filter and lambda
print(timeit.timeit('myidx , myel = list(filter(lambda el: el[1] > 0, enumerate([-1,-2,3,4,5,6])))[0]', number=100000))
0.0155314

statement = """
for idx, el in enumerate([-1,-2,3,4,5,6]):
    if el > 0:
        myidx, myel = idx, el
        break
"""

print(timeit.timeit(statement, number=100000))
0.04074070000000002

You can make use of the combination of lambda and filter like this:

mylist = [-1,-2,3,4,5,6]

myidx, myel = list(filter(lambda el: el[1] > 0, enumerate(mylist)))[0]
print("({}, {})".format(myidx, myel))

Explanation:

The filter() function which offers an elegant way to filter out all the elements takes in a function and a list as arguments. Here they are lambda and mylist. Since you want to get the corresponding index, we need to use enumerate to wrap up enumerate(mylist).

Basically, enumerate(mylist) returns a tuple of an index and the corresponding value. Our condition here is the comparison between the value and 0 so that's why we get el[1] instead of el[0] to compare with 0.

The results will be casted to list. This list includes all the pairs (index, value) that meet our condition. Here we want to get the first pair so that's why we have [0] at the end.

Output:

(2, 3)
Related