Iterate over a list based on list with set of iteration steps

This works:

l = [6,2,2,5,2,5,1,7,9,4]
w = [2,2,1,1]
k = 1

def take(xs, runs, skip_size):
    ixs = iter(xs)
    for run_size in runs:
        for _ in range(run_size ):
            yield next(ixs)
        for _ in range(skip_size):
            next(ixs)

result = list(take(l, w, k))
print(result)

Result:

[6, 2, 5, 2, 1, 9]

The function is what's called a generator, yielding one part of the result at a time, which is why it's combined into a list with list(take(l, w, k)).

Inside the function, the list xs that is passed in is wrapped in an iterator, to be able to take one item at a time with next().

runs defines how many items to take and yield, skip_size defines how many items to skip to skip after each 'run'.

As a bonus, here's a fun one-liner - if you can figure out why it works, I think you know enough about the problem to move on :)

[y for i, y in zip([x for xs in [[1] * aw + [0] * k for aw in w] for x in xs], l) if i]

@Grismar's implementation is excellent: straightforward, legible, and maintainable. Here is the compressed illegible version of the same:

from itertools import islice
from collections import deque

def take(xs, runs, skip_size):
    ixs = iter(xs)
    for run_size in runs:
        yield from islice(ixs, run_size)
        deque(islice(ixs, skip_size), maxlen=0)

The behavior is nearly identical in both cases.

v2

Based on @superb rain's fastest proposal, here is a slightly tweaked solution:

def take(xs, runs, skip_size):
    ixs = iter(xs)
    irs = iter(runs)
    yield from islice(ixs, next(irs, 0))
    for run in irs:
        yield from islice(ixs, skip_size, run + skip_size)

You can make a simple for loop and keep track of the current index where your range starts. Then in each iteration just update the start based on the previous one and your value of k.

l = [6,2,2,5,2,5,1,7,9,4]
w = [2,2,1,1]
k = 1

def get_slices(l, w, k):
    start = 0
    for n in w:
        yield from l[start: start+n]
        start += n + k
        

list(get_slices(l, w, k))
# [6, 2, 5, 2, 1, 9]

If you are using python > 3.8 you can stretch the readability a bit in exchange for brevity and fun with the walrus operator:

l = [6,2,2,5,2,5,1,7,9,4]
w = [2,2,1,1]
k = 1

start = -k

g = (slice(start:=start + k, start:=start + n) for n in w)
[j for slice in g for j in l[slice] ]
# [6, 2, 5, 2, 1, 9]

Tags:

Python