Repeat list elements based on values in another list

Viewed 259

I have two lists:

nums = [2, 3, 5]
mylist = ["aaa", "bbb", "ccc"]

What I am trying to achieve is repeat each element in mylist as many times as the corresponding values in nums. Expected output is:

["aaa", "aaa", "bbb", "bbb", "bbb", "ccc", "ccc", "ccc", "ccc", "ccc"]
4 Answers

Several solutions come to mind.

High level code with Counter

The easiest way is to use the elements() method for collections.Counter:

>>> from collections import Counter
>>> nums = [2, 3, 5]
>>> mylist = ["aaa", "bbb", "ccc"]
>>> list(Counter(dict(zip(mylist, nums))).elements())
['aaa', 'aaa', 'bbb', 'bbb', 'bbb', 'ccc', 'ccc', 'ccc', 'ccc', 'ccc']

Functional C-speed code with itertools

Or you can use the itertools module directly:

>>> from itertools import chain, repeat
>>> list(chain.from_iterable(map(repeat, mylist, nums)))
['aaa', 'aaa', 'bbb', 'bbb', 'bbb', 'ccc', 'ccc', 'ccc', 'ccc', 'ccc']

Native language syntax (no libraries)

To do this without libraries, a nested list comprehension would suffice:

>>> [elem for n, elem in zip(nums, mylist) for i in range(n)]
['aaa', 'aaa', 'bbb', 'bbb', 'bbb', 'ccc', 'ccc', 'ccc', 'ccc', 'ccc']

Let us do repeat

import pandas as pd
pd.Series(mylist).repeat(nums).tolist()
Out[554]: ['aaa', 'aaa', 'bbb', 'bbb', 'bbb', 'ccc', 'ccc', 'ccc', 'ccc', 'ccc']

Or try numpy

import numpy as np
np.repeat(mylist,nums)
Out[556]: 
array(['aaa', 'aaa', 'bbb', 'bbb', 'bbb', 'ccc', 'ccc', 'ccc', 'ccc',
       'ccc'], dtype='<U3')

Iterate both the lists in parallel using zip, then get the required value list with list multiplication by n:

nums = [2, 3, 5]
mylist = ["aaa", "bbb", "ccc"]
out = []
for v,n in zip(mylist, nums):
    out.extend([v]*n)

OUTPUT:

>>> out
['aaa', 'aaa', 'bbb', 'bbb', 'bbb', 'ccc', 'ccc', 'ccc', 'ccc', 'ccc']

Try:

nums = [2, 3, 5]
mylist = ["aaa", "bbb", "ccc"]

out = [s for n, s in zip(nums, mylist) for _ in range(n)]
print(out)

Prints:

['aaa', 'aaa', 'bbb', 'bbb', 'bbb', 'ccc', 'ccc', 'ccc', 'ccc', 'ccc']
Related