Home > other >  What is the pythonic way to skip to the last element of a list in a for loop?
What is the pythonic way to skip to the last element of a list in a for loop?

Time:08-31

I would like to perform an operation to each element in my list, and when a certain condition is met, to skip to the last element of the list.

Here is a MWE where I print all the items in a list until I reach my condition (item ==4), after which I manually repeat the print statement on the final element. The desired output is to print 0, 1, 2, 3, 4, 7:

my_list = [0, 1, 2, 3, 4, 5, 6, 7]
breaked_out = False
for item in my_list:
    print(item)
    if item == 4:
        breaked_out = True
        break
if breaked_out:
    print(my_list[-1])

I have this ugly use of a flag (breaked_out) and also need to repeat my print() command. This isn't particularly legible either.

I have a slightly better implementation in mind that uses a while loop instead:

my_list = [0, 1, 2, 3, 4, 5, 6, 7]
i = 0
while i < len(my_list):
    item = my_list[i]
    print(item)

    if item == 4:
        i = len(my_list)-1
    else:
        i  = 1

Here I'm not repeating my operation (in this case, print()) but I have to do this unpythonic index accounting.

Is there a more readable way to get this sort of iteration? Other things to add about my situation:

  • This is iterating on a list, not a generator, so I have access to len().
  • I need to loop through in this order, so I can't reverse(my_list) and treat the special case first.

CodePudding user response:

There is no Pythonic way to do exactly what you want

Not what you asked for, but less ugly

The minimal change, not so ugly solution that doesn't actually advance the iterator is to just put the code in the loop instead of having a flag variable:

my_list = [0, 1, 2, 3, 4, 5, 6, 7]

for item in my_list:
    print(item)
    if item == 4:
        print(my_list[-1])  # Handling of last element inlined
        break
# Optionally, an else: block can run to do something special when you didn't break
# which might be important if the item that's equal to 4 is the last or second
# to last item, where the former does the work for the final element twice,
# while the latter does it only once, but looks like it never found the element
# (processing all items without breaking looking the same as processing elements
# 0 through n - 1, then processing n separately, then breaking)
else:
    print("Processed them all!")

or to avoid processing the final element twice when it's the first element meeting the test criteria, use enumerate to track your position:

my_list = [0, 1, 2, 3, 4, 5, 6, 7]

for i, item in enumerate(my_list, 1):  # Lets us test against len(my_list) rather than len(my_list) - 1
    print(item)
    if item == 4 and i < len(my_list):  # Don't process last item if we just processed it!
        print(my_list[-1])  # Handling of last element inlined
        break

What you asked for:

There's only two ways I know of to do this, both of which involve converting to an iterator first so you can manipulate the iterator within the loop so it will skip to the last element for the next loop. In both cases, your original code changes to:

my_list = [0, 1, 2, 3, 4, 5, 6, 7]
lstiter = iter(my_list)
for item in lstiter:
    print(item)
    if item == 4:
        # iterator advance goes here

where that placeholder line at the bottom is what changes.

The documented, but slow approach

Using the consume recipe from itertools, advance it to near the end. You need to know where you are, so the for loop changes to:

for i, item in enumerate(lstiter, 1):  # Starting from 1 avoids needing an extra - 1 in
                                       # the length check and consume's second argument

and the placeholder is filled with:

if i < len(my_list):
    consume(lstiter, len(my_list) - i)

Downside: Advancing an arbitrary iterator manually like this is O(n) (it has to produce and discard all the values, which can take time for a large list).

The efficient, but undocumented approach

list iterators provide a __setstate__ method (it's used for pickleing them so they can be unpickled at the same offset in the underlying list). You can abuse this to change the position of the iterator to any place you like. For this, you keep the for loop without enumerate, and just fill the placeholder with:

lstiter.__setstate__(len(my_list) - 1)

which directly skips to the iterator such that the next element it produces will be the final element of the list. It's efficient, it's simple, but it's non-obvious, and I doubt any part of the spec requires that __setstate__ be provided at all, let alone implemented in this useful way (there are a bazillion methods you can choose from to implement pickling, and they could have selected another option). That said, the implementation is effectively required for all pickle protocol versions to date, for compatibility reasons (if they got rid of __setstate__, pickles produced on older Python would not be readable on modern Python), so it should be fairly reliable.

A warning:

If the final element of your list matches the condition, this will turn into an infinite loop, unlike the other solutions (break double processes the final element in that case, consume only processes each element at most once). breaking explicitly doesn't reenter the loop, so that's safe, and the consume recipe can't back up an iterator, so again, safe. But since this sets the position, it can set it back to the same position forever. If this is a possibility, I'd recommend the explicit break (using enumerate to check indices to avoid double-processing the final element), or failing that, you can add even more hackery by checking the length hint of the iterator to see if you were already at the end (and therefore should not adjust the position):

from operator import length_hint  # At top of file

my_list = [0, 1, 2, 3, 4, 5, 6, 7]
lstiter = iter(my_list)
for item in lstiter:
    print(item)
    if item == 4 and length_hint(lstiter) > 1:  # Don't set when finished or reaching last item anyway
        lstiter.__setstate__(len(my_list) - 1)

As an alternative to checking length_hint and similar hackery, you could use a flag variable that gets set to True when the condition passes and prevents reentering the if a second time, e.g.:

my_list = [0, 1, 2, 3, 4, 5, 6, 7]
lstiter = iter(my_list)
skipped = False
for i, item in enumerate(lstiter, 1):
    print(item)
    if item == 4 and not skipped and i < len(my_list):  # Don't set when finished or reaching last item anyway
        lstiter.__setstate__(len(my_list) - 1)
        skipped = True

but this is straying further and further from Pythonic with every change. :-)

CodePudding user response:

Don't do too much for a simple task

my_list = [0, 1, 2, 3, 4, 5, 6, 7]
for item in my_list:
    print(item)
    if item == 4:
        print(my_list[-1])
        break
    # task here
  • Related