Find Item After a Specific Item

Leave a reply

Here is a problem: given a sequence and an item in that sequence, find the item which follows.

Initial Solution

For those who code in C-like languages, the first attempt might looks like this:

def find_item_after(sequence, item):
    items_count = len(sequence)
    for index in range(items_count):
        if sequence[index] == item and index < items_count - 1:
            return sequence[index]
    return None

The problem with this solution is the ugliness of the code, and the inefficiency of indexing.

More Pythonic Solution

def find_item_after(sequence, item):
    for item_here, item_after in zip(sequence, sequence[1:]):
        if item_here == item:
            return item_after
    return None

This solution improves in readability, but at the expense of performance. First, the expression sequence[1:] creates a new sequence from the current one. Then, the zip function creates yet another sequence that are twice the size of the original requence. Note that in Python 3, zip will return a generator object instead of a sequence, which will help in term of efficiency.

Pythonic and Efficient

Until we use Python 3, we are better off using a generator version of zip, namely izip. We can also avoid the memory hogging aspect of sequence[1:] using islice:

from itertools import izip, islice
def find_item_after(sequence, item):
    for item_here, item_after in izip(sequence, islice(sequence, 1, None)):
        if item_here == item:
            return item_after
    return None

This solution is both Pythonic and efficient. However, we can still do better.

Using iter

def find_item_after(sequence, item):
    iterable = iter(sequence)
    for item_here in iterable:
        if item_here == item:
            return next(iterable, None)
    return None

This solution does not use any external library, nor does it create extra copy of the sequence. The next function takes a second parameter, the return value in case that we are at the end of the iterable, when next will generate a StopIteration exception.

What about Find Before?

We can adapt the solution which use iter for this purpose:

def find_item_before(sequence, item):
    iterable = iter(sequence)
    item_before = None
    for item_here in iterable:
        if item_here == item:
            return item_before
        item_before = item_here
    return None

However, the above is too clunky and not efficient: we have to assign a new value to item_before every time we go through the loop. What about the solution which uses izip and islice?

from itertools import izip, islice
def find_item_after(sequence, item):
    for item_before, item_here in izip(islice(sequence, 1, None), sequence):
        if item_here == item:
            return item_before
    return None

Note that in this solution, we start searching in the sequence starting with the second item (at index 1) and eliminate the first item from the search. If the caller wants an item before the first item, the loop will complete without any result and we return None.

Conclusion

This problem was originally an interview question, but it does have some practical application

Leave a comment