CODEWITHSUNDEEP
python
agf1997
agf1997

Reputation: 2898

Looping through a list with a custom counter

I have a list of directories in called dirs.

I also have a custom counter function called snake which uses the following syntax. 

snake(low, high, x)

An example of its usage would be 

[snake(1,3,x) for x in range(8)]
[1, 2, 3, 3, 2, 1, 1, 2]

dirs contains multiple identical directories. For example

[/my/path/one/, /my/path/one/, /my/path/one/, /my/path/one/, 
 /my/path/two/, /my/path/two/, /my/path/two/]

I'd like to use snake to create a variable with a set of indices associated with dirs where the snake pattern resets when a new directory name is encountered in the list. Something like:

ind = [1, 2, 3, 3, 1, 2, 3]

How can this be accomplished?

Upvotes: 0

Views: 70

Answers (1)

Mateen Ulhaq
Mateen Ulhaq

Reputation: 27271

In pure python:

def snake(low, high):
    return itertools.cycle(itertools.chain(
        range(low, high + 1),
        range(high, low - 1, -1)))

Usage:

>>> list(itertools.islice(snake(1, 3), 8))
[1, 2, 3, 3, 2, 1, 1, 2]

Let's now write a shy snake which panics and resets the generator when it encounters a new item:

def shy_snake(low, high, iterable):
    """d-d-doesss s-s-sssenpai noticesss me?"""
    prev = None
    for x in iterable:
        if x != prev:
            prev = x
            snake_it = snake(low, high)
        yield next(snake_it)

Actually, we can write this a bit more generically, to work with any sequence generator:

def shy(seq_func, iterable):
    prev = None
    for x in iterable:
        if x != prev:
            prev = x
            it = seq_func()
        yield next(it)

def shy_snake(low, high, iterable):
    """d-d-doesss s-s-sssenpai noticesss me?"""
    return shy(lambda: snake(low, high), iterable)

Testing:

>>> dirs = [1, 1, 1, 1, 2, 2, 2]
>>> list(shy_snake(1, 3, dirs))
[1, 2, 3, 3, 1, 2, 3]

Upvotes: 1

Related Questions