Middle-outwards alternating iterator

Question

I have a use-case where I want to iterate from the middle of an array outwards, so I generate a list of indices as follows:

import itertools


def array_indices_from_middle(array_len: int, num_indices: int = 16) -> list[int]:
    """
    Create indices for an array starting from the middle and then appending indices alternating the left and right.

    If more indices are needed than the array length, start again from the middle.
    """
    midpoint = (array_len - 1) // 2
    sel_inds = []

    left = range(midpoint - 1, -1, -1)
    right = range(midpoint + 1, array_len, 1)

    # right before left, because of how mid is chosen for even array lengths
    alt = [i for i in itertools.chain.from_iterable(itertools.zip_longest(right, left)) if i is not None]

    i = 0
    while len(sel_inds) < num_indices:
        if len(sel_inds) % array_len == 0:
            sel_inds.append(midpoint)
        else:
            sel_inds.append(alt[i % len(alt)])
            i += 1

    return sel_inds

It passes the following unit tests:

assert array_indices_from_middle(array_len=5, num_indices=5) == [2, 3, 1, 4, 0]
assert array_indices_from_middle(array_len=5, num_indices=4) == [2, 3, 1, 4]
assert array_indices_from_middle(array_len=5, num_indices=8) == [2, 3, 1, 4, 0, 2, 3, 1]

assert array_indices_from_middle(array_len=6, num_indices=6) == [2, 3, 1, 4, 0, 5]
assert array_indices_from_middle(array_len=6, num_indices=5) == [2, 3, 1, 4, 0]
assert array_indices_from_middle(array_len=6, num_indices=9) == [2, 3, 1, 4, 0, 5, 2, 3, 1]

My combination of iterators feels a bit hacky. Is there a simpler way of approaching this problem? I'm not worried about corner cases.

Answer 1

Your code is way too complicated, and not very efficient.

We can use an infinite iterator to repeatedly yield numbers back and forth.

We supply the infinite iterator with n, we yield n before the loop. We can then use an infinite loop, and yield n - 1 and n + 1 in the first iteration. We then yield n - 2 and n + 2 in the second iteration. In general, in ith iteration we yield n - i and n + i (one-based indexing). This will eventually yield all integers in both directions.

We can then just use itertools.islice to ask for the first length elements, length is the number of elements of the array. We can just create the iterator with length // 2 as its argument. length // 2 is the index of the middle element.

So your code becomes:

from itertools import count, islice
from typing import Generator, Iterable


def alternate(n: int = 0) -> Generator[int, None, None]:
    yield n
    for i in count(1):
        yield n - i
        yield n + i


def zigzag(arr: Iterable) -> list:
    l = len(arr)
    return [arr[i] for i in islice(alternate(l // 2), l)]

---

Here is another version, it is a bit more complicated, and therefore somewhat slower, but it doesn't utilize infinite iterators.

To achieve what you want, at iteration 0, the index change needs to be 0, then at iteration 1, the index should decrease by 1, at the second iteration the index needs to increase by 2 from the last index, and at iteration 3 the index needs to decrease by 3 from the last index...

The series of index changes is therefore 0, -1, 2, -3, 4, -5... It is like the series of natural numbers, but each odd number is negated.

So, here is the code:

def alternate_finite(n: int) -> Generator[int, None, None]:
    m = n // 2
    sign = 1
    for i in range(n):
        yield (m := m + i * sign)
        sign *= -1


def zigzag_finite(arr: Iterable) -> list:
    return [arr[i] for i in alternate_finite(len(arr))]

This version is considerably less efficient, but is cleverer and does more actual work unlike the lazy version.

In [395]: %timeit zigzag(range(4096))
846 µs ± 9.26 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

In [396]: %timeit zigzag_finite(range(4096))
1.14 ms ± 13.6 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Answer 2

I would be interested in what would be required to make this "scalable" in a numerical way, including using Numpy

As in the comments, depending on a lot of things:

• How often this is called
• How sensitive the application is to timing and memory
• How large the output array needs to be
• etc.

Numpy may be a more suitable choice than what you're doing, which - while it's Pythonic - might not perform like it needs to.

Consider instead:

import numpy as np


def array_indices_from_middle(array_len: int, num_indices: int) -> np.ndarray:
    non_repeating = np.empty(shape=array_len, dtype=int)
    non_repeating[0::2] = np.arange((array_len-1)//2, -1, -1)
    non_repeating[1::2] = np.arange((array_len-1)//2+1, array_len)

    out = np.tile(non_repeating, (num_indices + array_len - 1)//array_len)
    return out[:num_indices]


def test() -> None:
    assert np.all(array_indices_from_middle(array_len=5, num_indices=5) == [2, 3, 1, 4, 0])
    assert np.all(array_indices_from_middle(array_len=5, num_indices=4) == [2, 3, 1, 4])
    assert np.all(array_indices_from_middle(array_len=5, num_indices=8) == [2, 3, 1, 4, 0, 2, 3, 1])

    assert np.all(array_indices_from_middle(array_len=6, num_indices=6) == [2, 3, 1, 4, 0, 5])
    assert np.all(array_indices_from_middle(array_len=6, num_indices=5) == [2, 3, 1, 4, 0])
    assert np.all(array_indices_from_middle(array_len=6, num_indices=9) == [2, 3, 1, 4, 0, 5, 2, 3, 1])


if __name__ == '__main__':
    test()