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

Question

I want to iterate a given list based on a variable number of iterations stored in another list and a constant number of skips stored in as an integer.

Let's say I have 3 things -

1. l - a list that I need to iterate on (or filter)
2. w - a list that tells me how many items to iterate before taking a break
3. k - an integer that tells me how many elements to skip between each set of iterations.

To rephrase, w tells how many iterations to take, and after each set of iterations, k tells how many elements to skip.

So, if w = [4,3,1] and k = 2. Then on a given list (of length 14), I want to iterate the first 4 elements, then skip 2, then next 3 elements, then skip 2, then next 1 element, then skip 2.

Another example,

#Lets say this is my original list

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

Based on w and k, I want to iterate as -

6 -> Keep # w says keep 2 elements 
2 -> Keep
2 -> Skip # k says skip 1
5 -> Keep # w says keep 2 elements
2 -> Keep
5 -> Skip # k says skip 1
1 -> Keep # w says keep 1 element
7 -> Skip # k says skip 1
9 -> Keep # w says keep 1 element
4 -> Skip # k says skip 1

I tried finding something from itertools, numpy, a combination of nested loops, but I just can't seem to wrap my head around how to even iterate over this. Apologies for not providing any attempt, but I don't know where to start.

I dont necessarily need a full solution, just a few hints/suggestions would do.

Answer 1

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]

Answer 2

@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)

Answer 3

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]