How to make weighted list with list comprehension

Question

I have a list w where w[i] represents the number of times i should be present in the final result.

For example, if:

w = [1, 3, 4]

then the resulting list should be:

[0, 1, 1, 1, 2, 2, 2, 2]

Notice there is one 0, three 1's, and four 2's.

I have tried to accomplish this with the list comprehension:

w_list = [[i]*w[i] for i in range(len(w))]

but of course this doesn't quite work, giving me this:

[[0], [1, 1, 1], [2, 2, 2, 2]]

How can I write a list comprehension to get me the desired result, like in the example?

See also: https://treyhunner.com/2015/12/python-list-comprehensions-now-in-color/ — Karl Knechtel, Mar 08 '22 at 00:30

score 3 · Answer 1 · answered Mar 08 '22 at 00:22

3

You could use a double loop to flatten w_list:

w_list = [x for num, i in enumerate(w) for x in [num] * i]

or

w_list = [x for i in range(len(w)) for x in [i]*w[i]]

Output:

[0, 1, 1, 1, 2, 2, 2, 2]

answered Mar 08 '22 at 00:22

score 2 · Answer 2 · answered Mar 08 '22 at 00:32

2

Here's an alternative solution you might like.

w = [1, 3, 4]
print([i for i in range(len(w)) for _ in range(w[i])])

results in [0, 1, 1, 1, 2, 2, 2, 2] as desired.

Alternatively, [i for i,n in enumerate(w) for _ in range(n)] accomplishes the same thing.

answered Mar 08 '22 at 00:32

Ben Grossmann

2 Answers2