Using list comprehension in nested foor loops

mouad Et-tali 0
Viewed 34

I'm trying to refactor a code and i see 3 nested loops that can probably change to list comprehensions, but it's trickier than i thought. Can anyone help with this ?

l_ids=["0000000","0000010","0000020","0000030","0000040","0000050","0000060","0000070","0000080","0000090","0000100",
                 "0000110","0000120","0000130","0000140","0000150"]
l_years=[2018,2019,2020,2021,2022]
l_means=["car","train","plane","other","freight"]
l_id=[]
l_year=[]
l_mean=[]
l_energy=[]

for supp_id in  l_ids:
    for year in l_years:
        for mean in l_means:
            l_id.append(supp_id)
            l_year.append(year)
            l_mean.append(mean)
            l_energy.append(random.randint(100000000, 999999999))
nested-loops list-comprehension refactoring python
2 Answers

for l_id, l_year and l_mean you can use list comprehension, as below:

l_id = [supp_id for supp_id in l_ids for year in range(len(l_years)) for mean in range(len(l_means))]
l_year = [year for supp_id in range(len(l_ids)) for year in l_years for mean in range(len(l_means))]
l_mean = [mean for year in range(len(l_years)) for supp_id in range(len(l_ids)) for mean in l_means]

For l_energy, random has a .choices() function that you can use to get a list of random numbers, with the length specified. I have specified the length to be the product of the lengths of all 3 lists (.sample() could also be used, but will avoid repetition):

l_energy = random.choices(range(100000000, 999999999), k=len(l_ids)*len(l_years)*len(l_means))

But if you want that to be list comprehension, you can do:

l_energy = [random.randint(100000000, 999999999) for supp_id in range(len(l_ids)) for year in range(len(l_years)) for mean in range(len(l_means))]

I'm not sure what your final plan is for these lists, itertools may be a better solution.

Emi OB0

I am not sure why you would want to do that. I think, list comprehensions would make things more convoluted in this case, but, here is how I'd do it if I were to do it with list comprehensions:

l_ids = [str.format("{:0>7}", i) for i in range(0, 160, 10)]
l_years = list(range(2018, 2023))

rand_energy = lambda: random.randint(100000000, 999999999)

tmp = [(supp_id, year, mean, rand_energy())
        for supp_id in l_ids for year in l_years for mean in l_means]

l_id, l_year, l_mean, l_energy = [[t[i] for t in tmp]
                                      for i in range(len(tmp[0]))]
zerdeshtjuan0
Related