Akshay Sehgal
Akshay Sehgal

Reputation: 19332

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

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.

Upvotes: 9

Views: 1742

Answers (5)

superb rain
superb rain

Reputation: 5521

Similar to Mad Physicist's but using islice with start and stop argument:

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

A little benchmark with similar solutions:

2.63 μs  2.39 μs  2.29 μs  Grismar
2.11 μs  2.12 μs  2.11 μs  Mad_Physicist
1.49 μs  1.45 μs  1.35 μs  Mark_Meyer
1.34 μs  1.37 μs  1.38 μs  superb_rain
1.37 μs  1.40 μs  1.35 μs  superb_rain_2
1.34 μs  1.43 μs  1.35 μs  Mad_Physicist_2

Code:

from timeit import repeat
from itertools import islice
from collections import deque

def Grismar(xs, runs, skip):
    ixs = iter(xs)
    for run in runs:
        for _ in range(run ):
            yield next(ixs)
        for _ in range(skip):
            next(ixs)

def Mad_Physicist(xs, runs, skip):
    ixs = iter(xs)
    for run in runs:
        yield from islice(ixs, run)
        deque(islice(ixs, skip), 0)

def Mark_Meyer(xs, runs, skip):
    start = 0
    for run in runs:
        yield from xs[start : start+run]
        start += run + skip
        
def superb_rain(xs, runs, skip):
    ixs = iter(xs)
    start = 0
    for run in runs:
        yield from islice(ixs, start, start + run)
        start = skip

def superb_rain_2(xs, runs, skip):
    ixs = iter(xs)
    iruns = iter(runs)
    for run in iruns:
        yield from islice(ixs, run)
        for run in iruns:
            yield from islice(ixs, skip, skip + run)

def Mad_Physicist_2(xs, runs, skip):
    ixs = iter(xs)
    iruns = iter(runs)
    yield from islice(ixs, next(iruns, 0))
    for run in iruns:
        yield from islice(ixs, skip, run + skip)

funcs = Grismar, Mad_Physicist, Mark_Meyer, superb_rain, superb_rain_2, Mad_Physicist_2

l = [6,2,2,5,2,5,1,7,9,4]
w = [2,2,1,1]
k = 1
expect = [6, 2, 5, 2, 1, 9]
number = 10**5

for func in funcs:
    print(list(func(l, w, k)) == expect, func.__name__)
print()

tss = [[] for _ in funcs]
for _ in range(4):
    for func, ts in zip(funcs, tss):
        t = min(repeat(lambda: deque(func(l, w, k), 0), number=number)) / number
        ts.append(t)
        print(*('%.2f μs ' % (1e6 * t) for t in ts[1:]), func.__name__)
    print()

Upvotes: 1

Mad Physicist
Mad Physicist

Reputation: 114440

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

Upvotes: 4

Jonah.Checketts
Jonah.Checketts

Reputation: 64

Here is what I got:

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

for wItem in w:
    for i in range(wItem):
        testList.append(l[pos+i])
    pos += k + wItem
l = testList

test list will be a temporary list you store items in until the loop is completely executed.

I use the variable pos to keep track of the position you have looped to in the function and then ad the amount you want to skip after each iteration of the inner loop.

Upvotes: 0

Mark
Mark

Reputation: 92460

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]

Upvotes: 2

Grismar
Grismar

Reputation: 31354

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]

Upvotes: 6

Related Questions