CODEWITHSUNDEEP
python
Jeremy Allen
Jeremy Allen

Reputation: 6644

Skip n list items in for loop

I am building a list by looping over a source list. If I encounter a specific value I want to skip it and the next n items in the list.

Is there a nicer way to write the process_things function below?

def test_skip(thing):
    return thing == 'b'

def num_to_skip(thing):
    return 3

def process_things(things):
    """
    build list from things where we skip every test_skip() and the next num_to_skip() items
    """
    result = []
    skip_count = 0
    for thing in things:
        if test_skip(thing):
            skip_count = num_to_skip(thing)
        if skip_count > 0:
            skip_count -= 1
            continue
        result.append(thing)
    return result

source = list('abzzazyabyyab')
intended_result = list('aazyaa')

assert process_things(source) == intended_result

Upvotes: 1

Views: 346

Answers (1)

user2357112
user2357112

Reputation: 281758

You can use an iterator and the consume recipe from itertools. Calling consume(iterator, n) advances an iterator by n items. If you call it while for-looping over the iterator, the next iteration will start at the first item consume didn't consume.

import collections
from itertools import islice

def consume(iterator, n):
    "Advance the iterator n-steps ahead. If n is none, consume entirely."
    # Use functions that consume iterators at C speed.
    if n is None:
        # feed the entire iterator into a zero-length deque
        collections.deque(iterator, maxlen=0)
    else:
        # advance to the empty slice starting at position n
        next(islice(iterator, n, n), None)

def process_things(things):
    """
    build list from things where we skip every test_skip() and the next num_to_skip() items
    """
    result = []
    things = iter(things)
    for thing in things:
        if test_skip(thing):
            consume(things, num_to_skip(thing) - 1)
        else:
            result.append(thing)
    return result

Note that the way I've written process_things, when it hits an item that test_skip says to skip, it counts that item in the num_to_skip(thing) items to skip. That matches what your code does, but it doesn't quite match your description, which says "skip it and the next n items in the list".

Upvotes: 4

Related Questions