add an element after the first element which passes a condition in a list

user_16239119 0
Viewed 39
for i,k in enumerate(ls):
    if k == 3:
        ls.insert(i+1,"example")
        break

The above code iterates through a list ls and finds the first element that equals to 3 and inserts "example" after that element and stops. While the above code can be written as,

ls.insert(ls.index(3)+1,"example")

What is the most efficient way to write a program to enter a element after the first element that passes a condition such as,

    if k > 3:

or

    if isPrime(k):
list python loops performance
1 Answers

You could use an iterator for your condition and next:

ls = list(range(10))
# [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

idx = next((i for i in range(len(ls)) if ls[i]>3),  # could be isPrime(ls[i])
           len(ls)) # default insertion in the end
ls.insert(idx+1, 'X')
# [0, 1, 2, 3, 4, 'X', 5, 6, 7, 8, 9]

If you don't want to insert if the condition is not met:

idx = next((i for i in range(len(ls)) if ls[i]>10), None)
if idx is not None:
    ls.insert(idx+1, 'X')
# [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
alternative

The one-liner equivalent of your code (with the same flaws in case of non match) could use itertools.dropwhile (note the inverted condition):

ls = list(range(10))
from itertools import dropwhile

ls.insert(ls.index(next(dropwhile(lambda x: not x>3, ls)))+1, 'X')

Output: [0, 1, 2, 3, 4, 'X', 5, 6, 7, 8, 9]

user_163434640
Related