Bidirectional iterate over list in Python

Question

I wrote a short function to iterate over a list starting at a given index, then flipping back and forth between left and right values.

import itertools
from typing import Generator


def bidirectional_iterate(lst: list, index: int) -> Generator:
    for left, right in itertools.zip_longest(reversed(lst[:index]), lst[index:]):
        if right is not None:
            yield right
        if left is not None:
            yield left


arr = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
#      ^  ^  ^  ^  ^  ^  ^  ^  ^  ^
#      7  5  3  1  0  2  4  6  8  9   <- Expected order of iteration

print(list(bidirectional_iterate(arr, 4)))

Output:

[4, 3, 5, 2, 6, 1, 7, 0, 8, 9]

I feel like I've seen an algorithm like this before, but I couldn't find any reference to it.

This was the most Pythonic implementation I could come up with. The only improvement I can think to make to this implementation is a way to differentiate None list elements from the None values returned to pad itertools.zip_longest().

Are there any better/faster ways or any libraries that can accomplish this?

Answer 1

I like your code, but I think the biggest potential issue with it is that in each iteration of the loop, both left and right need to be checked to see if they're not None. As @FMc pointed out, another issue is that your reversed(lst[:index]) and lst[index:] lists both have to be built in memory before you even begin iterating. If you're dealing with very large lists, that might not be possible, and certainly won't be efficient.

Here's another alternative approach, that uses makes use of the roundrobin recipe at the end of the itertools docs. Rather than checking each iteration whether the generator has been exhausted, roundrobin removes an empty generator from the sequence of generators as soon as that generator raises StopIteration. (If you're happy with having a third-party import, you can just import it from more_itertools and it saves you a few lines of code). I also used generator expressions in place of the lists that you built in memory, and tinkered a little with your type hints to make them more expressive of what the code is achieving.

from itertools import cycle, islice
from typing import Iterator, Iterable, TypeVar, Sequence

T = TypeVar("T")

def roundrobin(*iterables: Iterable[T]) -> Iterator[T]:
    "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
    # Recipe credited to George Sakkis
    num_active = len(iterables)
    nexts = cycle(iter(it).__next__ for it in iterables)
    while num_active:
        try:
            for next in nexts:
                yield next()
        except StopIteration:
            # Remove the iterator we just exhausted from the cycle.
            num_active -= 1
            nexts = cycle(islice(nexts, num_active))


def bidirectional_iterate(xs: Sequence[T], index: int) -> Iterator[T]:
    left = (xs[i] for i in range((index - 1), -1, -1))
    right = (xs[i] for i in range(index, len(xs)))
    yield from roundrobin(right, left)

arr = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
print(list(bidirectional_iterate(arr, 4)))

Another option is to take away the function call and merge the two functions into one:

from itertools import cycle, islice
from typing import Iterator, TypeVar, Sequence

T = TypeVar("T")

def bidirectional_iterate(xs: Sequence[T], index: int) -> Iterator[T]:
    """Incorporates the 'roundrobin'  recipe from the itertools docs, credited to George Sakkis"""
    left = (xs[i] for i in range((index - 1), -1, -1))
    right = (xs[i] for i in range(index, len(xs)))
    num_active = 2
    nexts = cycle(iter(it).__next__ for it in (right, left))
    while num_active:
        try:
            for next in nexts:
                yield next()
        except StopIteration:
            # Remove the iterator we just exhausted from the cycle.
            num_active -= 1
            nexts = cycle(islice(nexts, num_active))

arr = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
print(list(bidirectional_iterate(arr, 4)))

Answer 2

Here's an alternative approach that requires no conditional logic and does not create intermediate sub-lists — although at the cost of using a third-party library:

from more_itertools import interleave_longest

def bidirectional_iterate(xs, index):
    left = range(index - 1, -1, -1)
    right = range(index, len(xs), 1)
    return [xs[i] for i in interleave_longest(right, left)]

Answer 3

Warning: this is probably not better - the index juggling is ugly - and it may or may not be faster; I haven't tested. It's possible for you to write this with no explicit loops or conditions (with sufficiently narrow definitions for loops and conditions) and no external libraries. You can assign to a list directly, abusing Python's slicing semantics:

from typing import TypeVar, Sequence, List

ItemT = TypeVar('ItemT')


def bidirectional_iterate(items: Sequence[ItemT], index: int) -> List[ItemT]:
    n = len(items)
    out = [None]*n

    # Increasing from index to end of items
    out[: 2*(n-index): 2] = items[index: index + n//2]

    # Increasing from beginning of items to index
    out[2*(n-index): : 2] = items[-n: index - n*3//2]

    # Decreasing from index to beginning of items
    out[
        min(-1, 2*index-n-1): : -2
    ] = items[index - n*3//2: index-n]

    # Decreasing from end of items to index
    out[n-1: 2*index: -2] = items[index + n//2:]

    return out


for i in range(10):
    print(bidirectional_iterate(tuple(range(10)), i))

If you end up using Numpy, this stands a good chance of executing more quickly than the iterative solution. I've tested this for its correctness at n=10 for all of the shown index values.

Answer 4

Your code is good from a readability standpoint. I appreciate the type hints. My only minor suggestions are using if __name__ == "__main__": for your driver code and lst instead of arr (arr usually refers to a NumPy array, but there are also builtin Python arrays that aren't lists).

Here's a somewhat terser suggestion that uses a generator expression and a classic iterable flattening pattern:

from itertools import zip_longest
from typing import Generator

def bidirectional_iterate(lst: list, index: int) -> Generator:
    zipped = zip_longest(lst[index:], reversed(lst[:index]))
    return (x for pair in zipped for x in pair if x is not None)

if __name__ == "__main__":
    lst = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
    #      ^  ^  ^  ^  ^  ^  ^  ^  ^  ^
    #      7  5  3  1  0  2  4  6  8  9   <- Expected order of iteration
    
    print(list(bidirectional_iterate(lst, 4)))

That said, it's an antipattern to return a linear-sized generator in an algorithm that needs to use linear space just to set up the generator. Returning a generator implies the algorithm uses space substantially less than the output set.

Changing the generator expression to a list comprehension gives:

def bidirectional_iterate(lst: list, index: int) -> Generator:
    zipped = zip_longest(lst[index:], reversed(lst[:index]))
    return [x for pair in zipped for x in pair if x is not None]

As a bonus, the caller need not use list() on the return value.

On the other hand, iteration algorithms are generally conducive to laziness, so it's usually a shame to throw the baby out with the bath water purely for the sake of more clever code; try to write it to keep the generator and get rid of the slices.

One approach that seems worth mentioning is treating the list as a graph with neighbors i - 1 if i < index else i + 1 running a breadth-first traversal on the list.

This incurs some overhead from the queue operations, so I wouldn't suggest it as categorically better than a purely index-based approach, but it's lazy, requires no dependencies and is pretty readable.

If q.pop(0) looks scary, keep in mind the list will never have more than 3 elements, so it's constant time, probably not much slower than a collections.deque.popleft(), which is the normal queue in Python. Feel free to use that and benchmark it if pop(0) (rightfully) makes you nervous.

def bidirectional_iterate(lst: list, index: int) -> Generator:
    q = [index, index - 1, index + 1]

    while q:
        i = q.pop(0)
        
        if i >= 0 and i < len(lst):
            yield lst[i]

            if i != index:
                q.append(i - 1 if i < index else i + 1)

Having offered this as a proof of concept, as suggested above, the nature of this algorithm is such that the entire list would need to be in memory anyway because the indices don't move in a forward-only direction.

Typical itertools algorithms and one-direction iteration functions like file readers seem like dramatically better fits for generators. For file readers, the input need not be fully in memory so your memory consumption can be line-by-line or even byte-by-byte.

In other cases, like itertools.permutations, the size of the result set is dramatically (exponentially) larger than the input iterables and liable to blow up memory. Other itertools functions like cycle return infinite sequences, another great use case for generators.

For this algorithm, I'd recommend going with the list version until you really find a critical use case for the generator. I'd be curious to hear what motivated you to pursue a generator in the first place.