Search code examples
pythonlistiterator

How to repeat list's elements with each different n times and in Python?


The idea is to repeat list's elements with each different n times as below.

ls = [7, 3, 11, 5, 2, 3, 4, 4, 2, 3]
id_list_fname = ['S11', 'S15', 'S16', 'S17', 'S19', 'S3', 'S4', 'S5', 'S6', 'S9']
all_ls = []
for id, repeat in zip(id_list_fname, ls):
    res = [ele for ele in[id] for i in range(repeat)]
    all_ls.append(res)

Consequently, I would like the result to be a single flat list, which I achieved as below.

def flatten(lst):
    for item in lst:
        if isinstance(item, list):
            yield from flatten(item)
        else:
            yield item

final_output = list(flatten(all_ls))

Output of the final_output:

['S11', 'S11', 'S11', 'S11', 'S11', 'S11', 'S11', 'S15', 'S15', 'S15',
 'S16', 'S16', 'S16', 'S16', 'S16', 'S16', 'S16', 'S16', 'S16', 'S16',
 'S16', 'S17', 'S17', 'S17', 'S17', 'S17', 'S19', 'S19', 'S3', 'S3',
 'S3', 'S4', 'S4', 'S4', 'S4', 'S5', 'S5', 'S5', 'S5', 'S6', 'S6',
 'S9', 'S9', 'S9']

But I wonder there exist much more compact approach or technique such as itertools to have the repeat elements as what I achieved using the code snippet above.


Solution

  • Numpy has the repeat function:

    np.repeat(id_list_fname, ls)
    

    If you absolutely need a list:

    np.repeat(id_list_fname, ls).tolist()
    

    If you don't want to use numpy, keep in mind that python allows nested comprehensions, which are essentially multiple for loops creating a single flat list. Your original loop can be written as a flat one-liner:

    [id for id, repeat in zip(id_list_fname, ls) for _ in range(repeat)]