Iterating over multiple lists sequentially with list index in Python-CodePudding

I'm looking for a library function, which iterates over multiple lists sequentially and produces their elements and index of the list the element is taken from. For example, this code:

a = [9, 1]
b = [3, 42]

for i, e in magic_chain(a, b):
  print(i, e)

would produce:

Does anything like this exist? I know how to write it, but I'm looking for existing library functionality.

CodePudding user response：

I'm not aware of anything that does exactly what you're looking for (at least, out of the functions in itertools). But, using chain() and repeat(), we can get something fairly concise that can be extended to an arbitrary number of iterables:

from itertools import chain, repeat

a = [9, 1]
b = [3, 42]

for i, e in chain.from_iterable(zip(repeat(idx), itr) for idx, itr in enumerate([a, b])):
  print(i, e)

CodePudding user response：

Not a single function, but we can combine a few to get a fast solution:

def magic_chain(*lists):
    return chain.from_iterable(map(zip, map(repeat, count()), lists))

A benchmark:

60.7 ms  magic_chain1
12.4 ms  magic_chain2

Code (Try it online!):

lists = [[0] * 1000] * 1000

def magic_chain1(*lists):
    for i, lst in enumerate(lists):
        for e in lst:
            yield i, e

def magic_chain2(*lists):
    return chain.from_iterable(map(zip, map(repeat, count()), lists))

from timeit import timeit
from collections import deque
from itertools import chain, repeat, count

funcs = magic_chain1, magic_chain2

for _ in range(3):
    for func in funcs:
        t = timeit(lambda: deque(func(*lists), 0), number=1)
        print('%4.1f ms ' % (t * 1e3), func.__name__)
    print()

CodePudding user response：

use enumerate:

a = [9, 1]
b = [3, 42]
for i, lst in enumerate([a, b]):
    for e in lst:
        print(i, e)